diff --git a/dlinfer/framework/lmdeploy_ext/cudagraph/ascend_cudagraph.py b/dlinfer/framework/lmdeploy_ext/cudagraph/ascend_cudagraph.py index baf75afa..6adbfada 100644 --- a/dlinfer/framework/lmdeploy_ext/cudagraph/ascend_cudagraph.py +++ b/dlinfer/framework/lmdeploy_ext/cudagraph/ascend_cudagraph.py @@ -19,6 +19,11 @@ BuffType = Dict[str, Tensor] +def _is_sparse_attention(graph_meta: CudaGraphMeta) -> bool: + """Whether this graph runs the device-only NSA/DSA attention path.""" + return getattr(graph_meta, "mla_index_topk", None) is not None + + # AscendCudaGraphMixin methods for cudagraph buffer management. def AscendCudaGraphMixin_support_cuda_graph( self, @@ -64,7 +69,13 @@ def AscendCudaGraphMixin_make_buffers_cudagraph( (max_batches, num_blocks), dtype=torch.int32, device=device ) - input_buffers["kv_seqlens"] = torch.ones(max_batches, dtype=torch.int32) + input_buffers["kv_seqlens"] = torch.ones( + max_batches, dtype=torch.int32, device=device + ) + if not _is_sparse_attention(graph_meta): + input_buffers["kv_seqlens_cpu"] = torch.ones( + max_batches, dtype=torch.int32 + ) input_buffers["kv_start_indices"] = -torch.ones( (max_tokens), dtype=torch.int32, device=device @@ -74,9 +85,11 @@ def AscendCudaGraphMixin_make_buffers_cudagraph( (max_tokens), dtype=torch.bool, device=device ) - input_buffers["attention_mask"] = torch.triu( - torch.ones(2048, 2048, dtype=torch.bool, device=device), diagonal=1 - ) + if not _is_sparse_attention(graph_meta): + input_buffers["attention_mask"] = torch.triu( + torch.ones(2048, 2048, dtype=torch.bool, device=device), + diagonal=1, + ) # ssm if graph_meta.is_ssm: @@ -87,21 +100,21 @@ def AscendCudaGraphMixin_make_buffers_cudagraph( max_batches, dtype=torch.int32, device=device ) - if max_batches != max_tokens: - max_q_seq_len = max_tokens // max_batches - input_buffers["q_seqlens"] = ( - torch.arange(1, max_batches + 1, dtype=torch.int32) * max_q_seq_len - ) - input_buffers["q_start_loc"] = ( - torch.arange(max_batches + 1, dtype=torch.int32, device=device) - * max_q_seq_len - ) - - else: - input_buffers["q_seqlens"] = torch.arange(1, max_batches + 1, dtype=torch.int32) - input_buffers["q_start_loc"] = torch.arange( - max_batches + 1, dtype=torch.int32, device=device + query_len = max_tokens // max_batches + input_buffers["q_seqlens"] = torch.full( + (max_batches,), query_len, dtype=torch.int32, device=device + ) + input_buffers["cu_seqlens_q"] = ( + torch.arange(max_batches + 1, dtype=torch.int32, device=device) + * query_len + ) + if not _is_sparse_attention(graph_meta): + input_buffers["cu_seqlens_q_cpu"] = ( + torch.arange(max_batches + 1, dtype=torch.int32) * query_len ) + # Keep q_start_loc as a compatibility view with its actual [B] start-offset + # semantics. Attention kernels consume cu_seqlens_q directly. + input_buffers["q_start_loc"] = input_buffers["cu_seqlens_q"][:-1] # mrope if graph_meta.use_mrope: @@ -128,7 +141,6 @@ def AscendCudaGraphMixin_fill_buffers_cudagraph( kv_start_indices: Tensor = attn_metadata.kv_start_indices moe_metadata = get_step_ctx_manager().current_context().moe_metadata x_active_mask: Tensor = moe_metadata.x_active_mask - q_start_loc: Tensor = attn_metadata.q_start_loc cache_seqlens: Tensor = attn_metadata.cache_seqlens is_multi_token_decoding = attn_metadata.is_multi_token_decoding @@ -154,6 +166,14 @@ def AscendCudaGraphMixin_fill_buffers_cudagraph( input_buffers["kv_seqlens"].fill_(0) input_buffers["kv_seqlens"][:batch_size] = kv_seqlens + if not _is_sparse_attention(graph_meta): + kv_seqlens_cpu: Tensor = attn_metadata.kv_seqlens_cpu + if kv_seqlens_cpu is None: + raise RuntimeError( + "Ascend paged-attention graph requires kv_seqlens_cpu" + ) + input_buffers["kv_seqlens_cpu"].fill_(0) + input_buffers["kv_seqlens_cpu"][:batch_size] = kv_seqlens_cpu input_buffers["kv_start_indices"].fill_(-1) input_buffers["kv_start_indices"][: kv_start_indices.size(0)] = kv_start_indices if x_active_mask is not None: @@ -171,7 +191,7 @@ def AscendCudaGraphMixin_fill_buffers_cudagraph( attn_metadata.cache_seqlens = input_buffers["cache_seqlens"] - if is_multi_token_decoding: + if is_multi_token_decoding and not _is_sparse_attention(graph_meta): attn_metadata.attention_mask = [input_buffers["attention_mask"]] if inputs_embeds is not None: @@ -185,10 +205,15 @@ def AscendCudaGraphMixin_fill_buffers_cudagraph( attn_metadata.block_offsets = input_buffers["block_offsets"] attn_metadata.kv_seqlens = input_buffers["kv_seqlens"] + attn_metadata.kv_seqlens_cpu = input_buffers.get("kv_seqlens_cpu") attn_metadata.kv_start_indices = input_buffers["kv_start_indices"] moe_metadata.x_active_mask = input_buffers["x_active_mask"] attn_metadata.q_start_loc = input_buffers["q_start_loc"] attn_metadata.q_seqlens = input_buffers["q_seqlens"] + attn_metadata.cu_seqlens_q = input_buffers["cu_seqlens_q"] + attn_metadata.cu_seqlens_q_cpu = input_buffers.get( + "cu_seqlens_q_cpu" + ) new_inputs = dict( past_key_values=past_key_values, @@ -223,7 +248,10 @@ def AscendCudaGraphMixin_update_context_cudagraph(self, graph_meta, context): input_buffers = graph_meta.input_buffers context.block_offsets = input_buffers["block_offsets"] context.kv_seqlens = input_buffers["kv_seqlens"] + context.kv_seqlens_cpu = input_buffers.get("kv_seqlens_cpu") context.q_start_loc = input_buffers["q_start_loc"] + context.q_seqlens = input_buffers["q_seqlens"] + context.cu_seqlens_q_cpu = input_buffers.get("cu_seqlens_q_cpu") context.kv_start_indices = input_buffers["kv_start_indices"] context.moe_metadata.x_active_mask = input_buffers["x_active_mask"] @@ -326,6 +354,7 @@ def __init__( input_buffers=dict(), output_buffers=dict(), vocab_size=self.model_config.vocab_size, + mla_index_topk=getattr(model_config, "mla_index_topk", None), is_ssm=len(model_config.states_shapes) > 0, use_mrope=model_config.use_mrope, ) @@ -380,15 +409,17 @@ def forward(self, **kwargs): context = self.ctx_mgr.current_context() self.model.update_context_cudagraph(self.meta, context) self._graph.replay() - if self.is_mla: - cpu_update_input = [ - {"actual_seq_kvlen": self.meta.input_buffers["kv_seqlens"].tolist()} - ] - else: - cpu_update_input = [ - {"actual_seq_lengths_kv": self.meta.input_buffers["kv_seqlens"]} - ] - self._graph.update(cpu_update_input=cpu_update_input) + if not _is_sparse_attention(self.meta): + if self.is_mla: + cpu_update_input = [ + {"actual_seq_kvlen": self.meta.input_buffers["kv_seqlens_cpu"]} + ] + else: + cpu_update_input = [ + {"actual_seq_lengths_kv": self.meta.input_buffers["kv_seqlens_cpu"]} + ] + self._graph.update(cpu_update_input=cpu_update_input) + # torch.npu.synchronize() output_buffers = self.meta.output_buffers output = self.model.get_outputs_cudagraph(output_buffers, **kwargs) return output @@ -471,7 +502,7 @@ def get_graph_key( if is_multi_token_decoding: q_seqlens = attn_metadata.q_seqlens - max_q_seq_len = attn_metadata.max_q_seq_len + max_q_seq_len = attn_metadata.max_q_seqlen batch_size = q_seqlens.size(0) if meta.padding_batch_size is None: new_batch_size = self._get_capture_tokens(batch_size) diff --git a/dlinfer/framework/lmdeploy_ext/device/__init__.py b/dlinfer/framework/lmdeploy_ext/device/__init__.py index 53fc28ea..19b77391 100644 --- a/dlinfer/framework/lmdeploy_ext/device/__init__.py +++ b/dlinfer/framework/lmdeploy_ext/device/__init__.py @@ -154,196 +154,762 @@ def _patched_rejection_sample( _reject_sampler_mod.rejection_sample = _patched_rejection_sample -##### patch cache engine ##### -def patch_contiguous_cache_engine(): - from lmdeploy.pytorch.config import CacheConfig, ModelConfig - from functools import reduce - from math import gcd - from lmdeploy.pytorch.engine import cache_engine +def patch_modelslim_quantization_config(): + """Add Ascend ModelSlim dispatch to LMDeploy's quantization config.""" + from collections.abc import Mapping + + from lmdeploy.pytorch.config import QuantizationConfig + + if getattr(QuantizationConfig, '_dlinfer_modelslim_patched', False): + return + + original_from_config = QuantizationConfig.from_config + original_get_quant_method = QuantizationConfig.get_quant_method @classmethod - def _cache_engine_allocate_caches( - cls, - num_blocks: int, - model_config: ModelConfig, - cache_config: CacheConfig, - world_size: int, - device: str, - ): - """Allocate caches.""" - num_layers = model_config.num_layers - - # get all descs - k_cache_desc = cls.get_k_cache_desc(model_config, cache_config, world_size) - v_cache_desc = cls.get_v_cache_desc(model_config, cache_config, world_size) - quant_cache_descs = cls.get_quant_cache_descs( - k_cache_desc, v_cache_desc, model_config, cache_config - ) - custom_cache_descs = cls.get_custom_cache_descs(model_config, cache_config) - cache_descs = ( - [k_cache_desc, v_cache_desc] + quant_cache_descs + custom_cache_descs + def custom_from_config(cls, hf_config): + quant_sources = [] + quant_config = getattr(hf_config, 'quantization_config', None) + if quant_config is not None: + quant_sources.append(quant_config) + for config_name in ('llm_config', 'text_config'): + nested_config = getattr(hf_config, config_name, None) + nested_quant_config = getattr(nested_config, + 'quantization_config', None) + if nested_quant_config is not None: + quant_sources.append(nested_quant_config) + + if not quant_sources: + return original_from_config(hf_config) + if any( + isinstance(config, Mapping) + and config.get('quant_method') == 'compressed-tensors' + for config in quant_sources): + return original_from_config(hf_config) + + quant_config = quant_sources[0] + if (not isinstance(quant_config, Mapping) + or quant_config.get('quant_method') != 'modelslim'): + return original_from_config(hf_config) + + quant_dtype = quant_config.get('quant_dtype') or 'int8' + resolved_quant_dtype = getattr(torch, quant_dtype, None) + if not isinstance(resolved_quant_dtype, torch.dtype): + raise ValueError( + f'Invalid quant dtype "{quant_dtype}" resolved from model ' + 'config; expected a torch.dtype attribute on torch.') + + ignored_layers = quant_config.get('ignored_layers', []) + if not ignored_layers: + ignored_layers = quant_config.get('modules_to_not_convert', []) + return cls( + quant_method='modelslim', + quant_dtype=resolved_quant_dtype, + scale_fmt=quant_config.get('scale_fmt'), + weight_block_size=quant_config.get('weight_block_size'), + activation_scheme=quant_config.get('activation_scheme'), + ignored_layers=ignored_layers, + fp8_quant_scope=quant_config.get('fp8_quant_scope'), + hf_quant_config=quant_config, ) - # get mempool size - mem_pool_size = 0 - alignments = [] - for desc in cache_descs: - mem_pool_size += desc.aligned_size - alignments.append(desc.alignment) - - # compute gcd of alignments - alignments_gcd = reduce(gcd, alignments) if alignments else 1 - assert ( - mem_pool_size % alignments_gcd == 0 - ), "mem_pool_size must be divisible by alignments_gcd" - - # create pool - mem_pool = torch.zeros( - (mem_pool_size // alignments_gcd, num_layers, num_blocks, alignments_gcd), - dtype=torch.uint8, - device=device, - ) + def get_modelslim_quant_method(self, prefix, module_kind): + if not prefix or module_kind == 'norm': + return None - # slice caches - caches = [] - remain_pool = mem_pool - for desc in cache_descs: - cache = ( - remain_pool[: desc.size // alignments_gcd, :, :, :] - .view(desc.dtype) - .view((num_layers, num_blocks, *desc.shape)) - ) - remain_pool = remain_pool[desc.aligned_size // alignments_gcd :, :, :, :] - caches.append(cache) - return mem_pool, caches + description = self.hf_quant_config.get('quant_description', {}) + if not description: + raise ValueError( + 'ModelSlim quantization requires quant_description metadata.') + + proj_name = prefix.rsplit('.', 1)[-1] + if module_kind == 'moe': + suffixes = ('0.gate_proj.weight', '0.up_proj.weight', + '0.down_proj.weight') + keys = [f'{prefix}.{suffix}' for suffix in suffixes] + elif proj_name == 'gate_up_proj': + parent = prefix.rsplit('.', 1)[0] + keys = [f'{parent}.gate_proj.weight', + f'{parent}.up_proj.weight'] + else: + keys = [f'{prefix}.weight'] - cache_engine.CacheEngine.allocate_caches = _cache_engine_allocate_caches + missing = [key for key in keys if key not in description] + if missing: + return None + quant_types = {description[key] for key in keys} + if len(quant_types) != 1: + raise ValueError( + f'ModelSlim fused module {prefix} mixes quant types: ' + f'{sorted(quant_types)}') + + quant_type = quant_types.pop() + if quant_type == 'FLOAT': + return None + if quant_type == 'W8A8_DYNAMIC': + return 'smooth_quant' + if quant_type == 'W8A8': + if module_kind == 'moe': + raise ValueError( + f'Static W8A8 MoE is not supported for {prefix}.') + return 'modelslim_w8a8_static' + raise ValueError( + f'Unsupported ModelSlim quant type {quant_type!r} for {prefix}.') + + def custom_get_quant_method(self, + prefix='', + module_kind='linear'): + if self.quant_method != 'modelslim': + return original_get_quant_method(self, prefix, module_kind) + if module_kind not in {'linear', 'moe', 'norm'}: + raise ValueError( + f'Unsupported quant module kind: {module_kind}') + return self._get_modelslim_quant_method(prefix, module_kind) + QuantizationConfig.from_config = custom_from_config + QuantizationConfig._get_modelslim_quant_method = ( + get_modelslim_quant_method) + QuantizationConfig.get_quant_method = custom_get_quant_method + QuantizationConfig._dlinfer_modelslim_patched = True -##### patch state cache engine ##### -def patch_state_cache_engine(): - from typing import List, Optional, Sequence, Tuple - from lmdeploy.pytorch.config import StateCacheSpec - from lmdeploy.pytorch.engine import cache_engine +def patch_deepseek_v32_config(): + """Allow the DeepSeek-V3.2 config builder to run on Ascend. - @staticmethod - def _state_cache_engine_allocate_caches( - num_caches: int, - state_shapes: List[Tuple[Tuple[int, ...], torch.dtype]], - device: torch.device, - state_specs: Optional[List[StateCacheSpec]] = None, - num_layers: Optional[int] = None, - ): - """Allocate cache implement. - - Each state is allocated as an independent contiguous tensor. A single - shared pool would give state views whose strides include the full pool - row, breaking NPU ops that require contiguous input. Layer-scoped named - caches use (num_rows, num_caches, *shape), matching lmdeploy's logical - layout while keeping every per-layer cache contiguous. - """ - - cache_dtype = torch.uint8 - state_specs = state_specs or [] - if (len(state_shapes) == 0 and len(state_specs) == 0) or num_caches == 0: - return torch.empty((0, 0), dtype=cache_dtype, device=device), [] - - resources = cache_engine.StateCacheEngine._get_state_cache_resources( - state_shapes, state_specs=state_specs, num_layers=num_layers + The upstream builder requires FlashMLA during config construction, while + Ascend uses dlinfer's Lightning Indexer instead. Temporarily report + FlashMLA as available while the upstream builder runs, then restore the + non-FlashMLA model semantics for the Ascend runtime. + """ + from lmdeploy.pytorch.configurations import deepseek_v2 as deepseek_v2_config + from lmdeploy.pytorch.configurations.deepseek_v32 import DeepseekV32ModelConfigBuilder + + if getattr(DeepseekV32ModelConfigBuilder, '_dlinfer_ascend_patched', False): + return + + original_build = DeepseekV32ModelConfigBuilder.build + original_flash_mla_available = deepseek_v2_config.flash_mla_available + + @classmethod + def custom_build(cls, hf_config, model_path: str | None = None, **kwargs): + device_type = kwargs.get('device_type', 'auto') + if device_type not in ('ascend', 'npu'): + return original_build(hf_config, model_path=model_path, **kwargs) + + deepseek_v2_config.flash_mla_available = lambda: True + try: + config = original_build(hf_config, model_path=model_path, **kwargs) + finally: + deepseek_v2_config.flash_mla_available = original_flash_mla_available + hf_config.use_flash_mla = False + + # Ascend uses dlinfer attention/indexer, not the CUDA FlashMLA path. + config.use_flash_mla = False + return config + + DeepseekV32ModelConfigBuilder.build = custom_build + DeepseekV32ModelConfigBuilder._dlinfer_ascend_patched = True + + +def patch_glm_moe_dsa_config(): + """Load Ascend ModelSlim metadata in the dlinfer configuration patch.""" + import json + import os + + from lmdeploy.pytorch.configurations.glm_moe_dsa import GlmMoeDsaModelConfigBuilder + from lmdeploy.utils import get_logger + + logger = get_logger('lmdeploy') + + if getattr(GlmMoeDsaModelConfigBuilder, '_dlinfer_modelslim_patched', False): + return + + original_build = GlmMoeDsaModelConfigBuilder.build + + @classmethod + def custom_build(cls, hf_config, model_path: str | None = None, **kwargs): + device_type = kwargs.get('device_type', 'auto') + modelslim_path = (os.path.join(model_path, 'quant_model_description.json') + if model_path else None) + if (device_type in ('ascend', 'npu') and modelslim_path + and os.path.isfile(modelslim_path)): + with open(modelslim_path, encoding='utf-8') as f: + quant_description = json.load(f) + if not isinstance(quant_description, dict): + raise TypeError(f'Expected a JSON object in {modelslim_path}.') + hf_config.quantization_config = { + 'quant_method': 'modelslim', + 'quant_dtype': 'int8', + 'quant_description': quant_description, + } + logger.info(f'Using Ascend ModelSlim quantization metadata from {modelslim_path}.') + return original_build(hf_config, model_path=model_path, **kwargs) + + GlmMoeDsaModelConfigBuilder.build = custom_build + GlmMoeDsaModelConfigBuilder._dlinfer_modelslim_patched = True + + +def patch_deepseek_v32_qkv(): + """Use separate Q-A and KV-A projections on Ascend. + + CUDA uses the merged ``fused_qkv_a_proj`` operator. ModelSlim W8A8 + checkpoints contain independent quantization metadata for ``q_a_proj`` + and ``kv_a_proj_with_mqa``, while Ascend does not provide the merged + operator. Keep this projection choice local to dlinfer. + """ + from lmdeploy.pytorch.models import deepseek_v32 + + attention_cls = deepseek_v32.DeepseekV32Attention + if getattr(attention_cls, '_dlinfer_ascend_qkv_patched', False): + return + + original_init = attention_cls.__init__ + + def custom_init(self, + config, + layer_idx, + dtype=None, + device=None, + all_reduce=True, + prefix=''): + if config.q_lora_rank is None: + return original_init(self, + config, + layer_idx, + dtype=dtype, + device=device, + all_reduce=all_reduce, + prefix=prefix) + + deepseek_v32.nn.Module.__init__(self) + self.layer_idx = layer_idx + quantization_config = getattr(config, 'quantization_config', None) + self.q_lora_rank = config.q_lora_rank + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.qk_rope_head_dim = config.qk_rope_head_dim + self.kv_lora_rank = config.kv_lora_rank + self.v_head_dim = config.v_head_dim + self.qk_nope_head_dim = config.qk_nope_head_dim + self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim + num_replicate_kv_heads = getattr(config, 'num_replicate_key_value_heads', 1) + num_key_value_heads = getattr(config, 'num_key_value_heads', 1) + use_flash_mla = getattr(config, 'use_flash_mla', False) + + self.q_a_proj = deepseek_v32.build_colwise_linear( + self.hidden_size, + config.q_lora_rank, + bias=config.attention_bias, + dtype=dtype, + device=device, + is_tp=False, + quant_config=quantization_config, + prefix=f'{prefix}.q_a_proj' if prefix else '', + ) + self.q_a_layernorm = deepseek_v32.RMSNorm( + config.q_lora_rank, + 1e-6, + quant_config=quantization_config, + dtype=deepseek_v32.torch.float32, + device=device, + ) + self.q_b_proj = deepseek_v32.build_colwise_linear( + config.q_lora_rank, + self.num_heads * self.q_head_dim, + bias=False, + dtype=dtype, + device=device, + is_tp=True, + quant_config=quantization_config, + prefix=f'{prefix}.q_b_proj' if prefix else '', + ) + self.kv_a_proj_with_mqa = deepseek_v32.build_colwise_linear( + self.hidden_size, + config.kv_lora_rank + config.qk_rope_head_dim, + bias=config.attention_bias, + dtype=dtype, + device=device, + is_tp=False, + quant_config=quantization_config, + prefix=f'{prefix}.kv_a_proj_with_mqa' if prefix else '', + ) + self.kv_a_layernorm = deepseek_v32.RMSNorm( + config.kv_lora_rank, + 1e-6, + quant_config=quantization_config, + dtype=deepseek_v32.torch.float32, + device=device, ) + self.kv_b_proj = deepseek_v32.build_colwise_linear( + config.kv_lora_rank, + self.num_heads * (config.qk_nope_head_dim + self.v_head_dim), + bias=False, + dtype=dtype, + device=device, + is_tp=True, + quant_config=quantization_config, + prefix=f'{prefix}.kv_b_proj' if prefix else '', + ) + self.kc = deepseek_v32.DeepseekV2BMM( + self.num_heads, + config.qk_nope_head_dim, + config.kv_lora_rank, + dtype=dtype, + device=device, + ) + self.apply_rotary_pos_emb = deepseek_v32.ApplyRotaryEmb() + self.softmax_scale = self.q_head_dim**-0.5 + + rope_scaling = deepseek_v32.get_rope_parameters(config) + if rope_scaling is not None: + mscale_all_dim = rope_scaling.get('mscale_all_dim', 0) + if mscale_all_dim: + scaling_factor = rope_scaling['factor'] + mscale = deepseek_v32.yarn_get_mscale(scaling_factor, mscale_all_dim) + self.softmax_scale = self.softmax_scale * mscale * mscale + + self.attn_fwd = deepseek_v32.Attention( + self.num_heads, + config.kv_lora_rank + self.qk_rope_head_dim, + scale=self.softmax_scale, + num_kv_heads=num_key_value_heads, + v_head_size=config.kv_lora_rank, + num_replicate_kv_heads=num_replicate_kv_heads, + use_flash_mla=use_flash_mla, + mla_index_topk=config.index_topk, + ) + self.vc = deepseek_v32.DeepseekV2BMM( + self.num_heads, + config.kv_lora_rank, + self.v_head_dim, + dtype=dtype, + device=device, + ) + self.o_proj = deepseek_v32.build_o_proj( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=config.attention_bias, + dtype=dtype, + device=device, + is_tp=True, + quant_config=quantization_config, + all_reduce=all_reduce, + prefix=f'{prefix}.o_proj' if prefix else '', + ) + self.indexer = self._build_indexer(config, layer_idx, dtype, device, prefix) + + def custom_qkv_proj(self, hidden_states, num_heads): + nope_size = self.kv_lora_rank + pe_size = self.qk_rope_head_dim + if self.q_lora_rank is None: + q_a_states = hidden_states + key_states = self.kv_a_proj_with_mqa(hidden_states[0, :, None]) + else: + q_a_states = self.q_a_proj(hidden_states) + key_states = self.kv_a_proj_with_mqa(hidden_states[0, :, None]) + + query_states, q_pe, qr = self._q_proj(q_a_states, num_heads, nope_size, pe_size) + key_states, value_states, k_pe = self._kv_proj(key_states, nope_size) + return query_states, key_states, value_states, q_pe, k_pe, qr - # Allocate each state as a separate contiguous tensor. - caches = [] - for resource in resources: - desc = resource.desc - cache_shape = (num_caches, *desc.shape) - if resource.layout is not None: - cache_shape = (resource.num_rows, num_caches, *desc.shape[1:]) - cache = torch.zeros(cache_shape, dtype=desc.dtype, device=device) - caches.append(cache) - - # mem_pool is used by two callers: - # 1. get_cache_state_size(): always calls with device='meta' to compute byte - # counts — the tensor is never materialised on a real device. - # 2. init_caches()/copy_caches(): patched below to operate on the independent - # cache tensors directly, so they no longer touch mem_pool at all. - # Therefore we only need a correctly-sized pool on 'meta'; for real devices we - # return an empty placeholder to avoid doubling the state-cache memory footprint. - total_bytes = sum(resource.desc.aligned_size for resource in resources) - if str(device) == "meta": - mem_pool = torch.empty( - (num_caches, total_bytes), dtype=cache_dtype, device=device + attention_cls.__init__ = custom_init + attention_cls._qkv_proj = custom_qkv_proj + attention_cls._dlinfer_ascend_qkv_patched = True + + +def patch_glm_moe_dsa_norm_dtype(): + """Cast GLM normalization layers to the model dtype on Ascend.""" + from lmdeploy.pytorch.models.glm_moe_dsa import ( + GlmMoeDsaDecoderLayer, + GlmMoeDsaModel, + ) + + if getattr(GlmMoeDsaModel, '_dlinfer_norm_dtype_patched', False): + return + + original_decoder_init = GlmMoeDsaDecoderLayer.__init__ + original_model_init = GlmMoeDsaModel.__init__ + + def custom_decoder_init(self, + config, + layer_idx, + dtype=None, + device=None, + prefix=''): + original_decoder_init(self, + config, + layer_idx, + dtype=dtype, + device=device, + prefix=prefix) + if dtype is not None: + self.input_layernorm.to(dtype=dtype) + self.post_attention_layernorm.to(dtype=dtype) + + def custom_model_init(self, config, dtype=None, device=None): + original_model_init(self, config, dtype=dtype, device=device) + if dtype is not None: + self.norm.to(dtype=dtype) + + GlmMoeDsaDecoderLayer.__init__ = custom_decoder_init + GlmMoeDsaModel.__init__ = custom_model_init + GlmMoeDsaModel._dlinfer_norm_dtype_patched = True + + +def patch_deepseek_v2_moe(): + """Use the Ascend EP reduction semantics for DeepSeek MoE.""" + from lmdeploy.pytorch.models import deepseek_v2 + + moe_cls = deepseek_v2.DeepseekV2MoE + if getattr(moe_cls, '_dlinfer_ascend_moe_patched', False): + return + + original_init = moe_cls.__init__ + + def custom_init(self, + config, + layer_idx, + dtype=None, + device=None, + all_reduce=True, + prefix=''): + original_init(self, + config, + layer_idx, + dtype=dtype, + device=device, + all_reduce=all_reduce, + prefix=prefix) + + dist_ctx = deepseek_v2.get_dist_manager().current_context() + dist_config = dist_ctx.dist_config + self._all_reduce = (all_reduce and dist_config.dp == 1 + and dist_config.world_size > 1 + and dist_config.ep == 1) + self._all_reduce_shared_experts = ( + all_reduce and dist_config.dp == 1 and dist_config.ep > 1 + and dist_config.mlp_tp > 1) + self._shared_expert_tp_group = None + if self._all_reduce_shared_experts: + self._shared_expert_tp_group = dist_ctx.mlp_tp_group.gpu_group + + def custom_forward(self, hidden_states, all_routed_experts=None): + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + routed_experts = None + if all_routed_experts is not None: + routed_experts = all_routed_experts[:, self.layer_idx, :] + topk_weights, topk_ids = self.gate( + hidden_states, routed_experts=routed_experts) + + out_states = self.experts(hidden_states, topk_weights, topk_ids) + if self.shared_experts is not None: + shared_states = self.shared_experts(hidden_states) + # EP already combines routed expert outputs. Only the shared expert + # output remains sharded over the MLP TP group. + if self._all_reduce_shared_experts: + deepseek_v2.dist.all_reduce( + shared_states, group=self._shared_expert_tp_group) + out_states += shared_states + out_states = out_states.reshape(batch_size, sequence_length, -1) + + if self._all_reduce: + deepseek_v2.dist.all_reduce(out_states) + return out_states + + moe_cls.__init__ = custom_init + moe_cls.forward = custom_forward + moe_cls._dlinfer_ascend_moe_patched = True + + +def patch_deepseek_v2_modelslim_weight_loader(): + """Adapt ModelSlim auxiliary checkpoint tensors on Ascend.""" + from lmdeploy.pytorch.models import deepseek_v2 + + model_cls = deepseek_v2.DeepseekV2ForCausalLM + if getattr(model_cls, '_dlinfer_modelslim_weight_loader_patched', False): + return + + original_load_weight_attention = model_cls._load_weight_attention + original_load_weights = model_cls.load_weights + + def map_modelslim_param_name(self, name, params_dict): + quantization_config = getattr(self.config, + 'quantization_config', None) or {} + if quantization_config.get('quant_method') != 'modelslim': + return name + if name.endswith('.weight_offset'): + return None + if name.endswith('.weight_scale'): + mapped_name = name.removesuffix('.weight_scale') + '.scale' + return mapped_name if mapped_name in params_dict else None + return name + + def custom_load_weight_experts(self, name, loaded_weight, params_dict, + expert_params_mapping): + for param_name, weight_name, expert_id, shard_id in expert_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + name = self._map_modelslim_param_name(name, params_dict) + if name is None: + return + param = params_dict[name] + deepseek_v2.load_weight( + param, + loaded_weight, + expert_id=expert_id, + shard_id=shard_id, ) + break else: - mem_pool = torch.empty(0, dtype=cache_dtype, device=device) - return mem_pool, caches - - def _state_cache_slot_dim(self, cache_idx: int): - """Return the state-slot dimension for one cache tensor.""" - cache_names = getattr(self, "_state_cache_names", []) - layer_maps = getattr(self, "_state_cache_layer_maps", {}) - if cache_idx < len(cache_names) and cache_names[cache_idx] in layer_maps: - return 1 - return 0 - - def _state_cache_engine_init_caches(self, idx: torch.Tensor, mask: torch.Tensor): - """Initialize state caches by zeroing each individual cache tensor.""" - if idx is None: - return - if len(self._state_caches) <= 0: + name = self._map_modelslim_param_name(name, params_dict) + if name is None: + return + deepseek_v2.load_weight(params_dict[name], loaded_weight) + + def custom_load_weight_attention(self, name, loaded_weight, params_dict, + update_pe_mapping): + mapped_name = self._map_modelslim_param_name(name, params_dict) + if mapped_name is None: return - num_caches = self.cache_config.num_state_caches - cache_masks = torch.zeros((num_caches,), dtype=torch.bool, device=idx.device) - cache_masks.index_copy_(0, idx, mask) - for cache_idx, cache in enumerate(self._state_caches): - slot_dim = _state_cache_slot_dim(self, cache_idx) - mask_shape = [1] * cache.dim() - mask_shape[slot_dim] = num_caches - reshaped_mask = cache_masks.view(mask_shape) - cache.masked_fill_(reshaped_mask, 0) - - def _state_cache_engine_copy_caches( - self, src_idx: int | Sequence[int], dst_idx: int | Sequence[int] - ): - """Copy slots between independently allocated state-cache tensors.""" - if len(self._state_caches) <= 0: + # Input quantization metadata is shared by the whole projection and + # has shape [1]. It must not enter DeepSeek's output-channel RoPE + # permutation, which expects dim 0 to be divisible by head_dim. + if mapped_name.endswith(('.input_scale', '.input_offset')): + deepseek_v2.load_weight(params_dict[mapped_name], loaded_weight) return + # Delegate output-channel metadata (for example a mapped dynamic + # weight scale) using its actual parameter name so that it receives + # the same RoPE permutation as the corresponding projection weight. + return original_load_weight_attention( + self, + mapped_name, + loaded_weight, + params_dict, + update_pe_mapping, + ) - src_list = self._index_list(src_idx) - dst_list = self._index_list(dst_idx) - if len(src_list) != len(dst_list): - raise ValueError( - "src_idx and dst_idx must have the same number of elements." - ) - if len(src_list) == 0: - return + def custom_load_weights(self, weights): + quantization_config = getattr(self.config, + 'quantization_config', None) or {} + if quantization_config.get('quant_method') != 'modelslim': + return original_load_weights(self, weights) + + params_dict = dict(self.named_parameters()) + stacked_params_mapping = [ + ('.gate_up_proj', '.gate_proj'), + ('.gate_up_proj', '.up_proj'), + ] + if not getattr(self.config, 'use_mla', True): + stacked_params_mapping.extend([ + ('.qkv_proj', '.q_proj'), + ('.qkv_proj', '.k_proj'), + ('.qkv_proj', '.v_proj'), + ]) + + def convert_weights(): + for name, loaded_weight in weights: + is_attention = ('.self_attn' in name + and getattr(self.config, 'use_mla', True)) + if '.experts' in name or is_attention: + yield name, loaded_weight + continue + if name.endswith('.weight_offset'): + continue + if name.endswith('.weight_scale'): + mapped_name = name.removesuffix('.weight_scale') + '.scale' + param_name = mapped_name + for fused_name, shard_name in stacked_params_mapping: + if shard_name in param_name: + param_name = param_name.replace(shard_name, + fused_name) + break + if param_name not in params_dict: + continue + name = mapped_name + yield name, loaded_weight + + return original_load_weights(self, convert_weights()) + + model_cls._map_modelslim_param_name = map_modelslim_param_name + model_cls._load_weight_experts = custom_load_weight_experts + model_cls._load_weight_attention = custom_load_weight_attention + model_cls.load_weights = custom_load_weights + model_cls._dlinfer_modelslim_weight_loader_patched = True + + +def patch_glm_moe_dsa_weight_loader(): + """Ignore the ModelSlim QuaRot-only MTP weight on Ascend.""" + from lmdeploy.pytorch.models.glm_moe_dsa import GlmMoeDsaForCausalLM + + if getattr(GlmMoeDsaForCausalLM, + '_dlinfer_ascend_weight_loader_patched', False): + return + + original_load_weights = GlmMoeDsaForCausalLM.load_weights + + def custom_load_weights(self, weights): + weights = ((name, weight) for name, weight in weights + if name != 'rot.weight') + return original_load_weights(self, weights) + + GlmMoeDsaForCausalLM.load_weights = custom_load_weights + GlmMoeDsaForCausalLM._dlinfer_ascend_weight_loader_patched = True + + +def patch_glm_moe_dsa_indexer(): + """Use the Ascend unfused, non-Hadamard DSA indexer path. + + The common GLM implementation keeps CUDA's fused projection and + Hadamard preprocessing semantics. Ascend's Lightning Indexer consumes + the separate BF16 projection outputs directly, so this adaptation stays + local to the dlinfer device patch. + """ + from lmdeploy.pytorch import envs as lmdeploy_envs + from lmdeploy.pytorch.models.glm_moe_dsa import GlmMoeDsaIndexer + + if getattr(GlmMoeDsaIndexer, '_dlinfer_ascend_patched', False): + return + + original_init = GlmMoeDsaIndexer.__init__ + + def custom_init(self, + config, + layer_idx, + dtype=None, + device=None, + prefix=''): + # Force the common constructor to materialize wk and weights_proj. + original_disable = lmdeploy_envs.disable_dsa_indexer_fusion + lmdeploy_envs.disable_dsa_indexer_fusion = True + try: + original_init(self, + config, + layer_idx, + dtype=dtype, + device=device, + prefix=prefix) + finally: + lmdeploy_envs.disable_dsa_indexer_fusion = original_disable + self.use_fusion = False + + def custom_forward(self, x, qr, freqs_cis, attn_metadata=None): + # This is the common unfused path without CUDA-only Hadamard rotation. + q = self.wq_b(qr).unflatten(-1, (-1, self.head_dim)) + q_pe, q_nope = torch.split( + q, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1) + k = self.k_norm(self.wk(x)) + k_pe, k_nope = torch.split( + k, [self.rope_head_dim, self.head_dim - self.rope_head_dim], dim=-1) + q_pe, k_pe = self._apply_rotary_pos_emb(q_pe, k_pe, freqs_cis) + q = torch.cat([q_pe, q_nope], dim=-1) + k = torch.cat([k_pe[0], k_nope[0, :, None]], dim=-1) + weights = self.weights_proj(x) * self.n_heads**-0.5 + return self.indexer_topk(q[0], + k[:, 0], + weights[0], + attn_metadata=attn_metadata) + + GlmMoeDsaIndexer.__init__ = custom_init + GlmMoeDsaIndexer.forward = custom_forward + GlmMoeDsaIndexer._dlinfer_ascend_patched = True - num_caches = self.cache_config.num_state_caches - self._validate_index_bounds(src_list, num_caches) - self._validate_index_bounds(dst_list, num_caches) - dst_set = set(dst_list) - if len(dst_set) != len(dst_list): - raise ValueError("dst_idx must not contain duplicate entries.") - if not set(src_list).isdisjoint(dst_set): - raise ValueError( - "src_idx and dst_idx must not overlap for stream-ordered state copies." - ) - for cache_idx, cache in enumerate(self._state_caches): - slot_dim = _state_cache_slot_dim(self, cache_idx) - for src, dst, length in self._copy_ranges(src_list, dst_list): - src_slice = [slice(None)] * cache.dim() - dst_slice = [slice(None)] * cache.dim() - src_slice[slot_dim] = slice(src, src + length) - dst_slice[slot_dim] = slice(dst, dst + length) - cache[tuple(dst_slice)].copy_( - cache[tuple(src_slice)], non_blocking=True - ) +##### patch cache engine ##### + + +def patch_glm_moe_dsa_split_cache(): + """Use independent contiguous noPE and RoPE caches for Ascend DSA.""" + from lmdeploy.pytorch.configurations.glm_moe_dsa import ( + GlmMoeDsaModelConfigBuilder, + ) + from lmdeploy.pytorch.distributed import get_dist_manager + from lmdeploy.pytorch.models.glm_moe_dsa import GlmMoeDsaAttention + + if getattr(GlmMoeDsaModelConfigBuilder, + '_dlinfer_split_cache_patched', False): + return + + original_build = GlmMoeDsaModelConfigBuilder.build + + @classmethod + def custom_build(cls, + hf_config, + model_path: str | None = None, + **kwargs): + config = original_build(hf_config, + model_path=model_path, + **kwargs) + # Cache only the RoPE key in K and the latent/noPE value in V. Their + # combined width is unchanged, but each cache can now be contiguous. + config.k_head_dim = hf_config.qk_rope_head_dim + config.v_head_dim = hf_config.kv_lora_rank + config.split_mla_kv_cache = True + return config + + def custom_forward( + self, + hidden_states, + rotary_pos_emb, + past_key_value=None, + attn_metadata=None, + topk_indices_buffer=None, + skip_topk: bool = False, + ): + dist_config = get_dist_manager().current_config() + num_heads = (self.num_heads if dist_config.dp > 1 else + self.num_heads // dist_config.attn_tp) + nope_size = self.kv_lora_rank + q_len = hidden_states.size(1) + + query_states, key_states, value_states, q_pe, k_pe, qr = ( + self._qkv_proj(hidden_states, num_heads=num_heads)) + cos, sin = rotary_pos_emb + q_pe, k_pe = self.apply_rotary_pos_emb(q_pe, + k_pe, + cos, + sin, + inplace=False) + query_states[..., nope_size:] = q_pe + key_states[..., nope_size:] = k_pe + + if topk_indices_buffer is None: + raise RuntimeError( + f'Layer {self.layer_idx} requires a DSA top-k indices buffer.') + if self.indexer is not None and not skip_topk: + topk_indices = topk_indices_buffer.write( + self.indexer(hidden_states, + qr, + rotary_pos_emb, + attn_metadata=attn_metadata)) + else: + topk_indices = topk_indices_buffer.read(q_len, + hidden_states.device) + + rope_cache, nope_cache = past_key_value[:2] + + attn_output = self.attn_fwd( + query_states, + key_states, + value_states, + rope_cache, + nope_cache, + attn_metadata, + k_scales_zeros=(None if len(past_key_value) == 2 else + past_key_value[2]), + v_scales_zeros=(None if len(past_key_value) == 2 else + past_key_value[3]), + nsa_indices=topk_indices, + ) + attn_bmm_out = attn_output.new_empty(q_len, num_heads, + self.v_head_dim) + self.vc(attn_output, attn_bmm_out) + return self.o_proj(attn_bmm_out.flatten(-2, -1)[None]) - cache_engine.StateCacheEngine.allocate_caches = _state_cache_engine_allocate_caches - cache_engine.StateCacheEngine.init_caches = _state_cache_engine_init_caches - cache_engine.StateCacheEngine.copy_caches = _state_cache_engine_copy_caches + GlmMoeDsaModelConfigBuilder.build = custom_build + GlmMoeDsaModelConfigBuilder._dlinfer_split_cache_patched = True + GlmMoeDsaAttention.forward = custom_forward def patch_gated_delta_net(): @@ -366,9 +932,9 @@ def __init__( attn_metadata: Any, ): self.is_decoding = attn_metadata.is_decoding - self.cu_seqlens = attn_metadata.q_start_loc + self.cu_seqlens = attn_metadata.cu_seqlens_q self.is_multi_token_decoding = attn_metadata.is_multi_token_decoding - self.max_q_seq_len = attn_metadata.max_q_seq_len + self.max_q_seq_len = attn_metadata.max_q_seqlen self.num_spec_tokens = get_step_ctx_manager().build_ctx.num_spec_tokens self.cache_seqlens = getattr(attn_metadata, "cache_seqlens", None) @@ -951,11 +1517,18 @@ def vendor_device_init(): import_vendor_module(vendor_name) patch_compiled_func() patch_async_sampling_logits() - if vendor_name in ["camb", "ascend"]: - patch_contiguous_cache_engine() if vendor_name == "ascend": patch_rejection_sampler() - patch_state_cache_engine() + patch_modelslim_quantization_config() + patch_glm_moe_dsa_config() + patch_deepseek_v32_config() + patch_deepseek_v32_qkv() + patch_glm_moe_dsa_norm_dtype() + patch_deepseek_v2_moe() + patch_deepseek_v2_modelslim_weight_loader() + patch_glm_moe_dsa_weight_loader() + patch_glm_moe_dsa_indexer() + patch_glm_moe_dsa_split_cache() patch_gated_delta_net() patch_qwen3_5() diff --git a/dlinfer/ops/llm.py b/dlinfer/ops/llm.py index af2d8db1..57f2efef 100644 --- a/dlinfer/ops/llm.py +++ b/dlinfer/ops/llm.py @@ -16,6 +16,9 @@ __all__ = [ "add_rms_norm", "apply_rotary_pos_emb", + "apply_rotary_pos_emb_interleaved", + "lightning_indexer", + "sparse_flash_attention", "prefill_attention", "incre_flash_attention", "fill_kv_cache", @@ -29,15 +32,48 @@ "get_cache_len", "weight_quant_matmul", "fused_moe", + "fused_moe_w8a8", "linear", "dynamic_quant", "linear_w8a8", + "linear_w8a8_static", "rms_norm_w8a8", "add_rms_norm_w8a8", "transdata", ] +def _lightning_indexer_abstract( + query, + key, + weights, + actual_seq_lengths_query, + actual_seq_lengths_key, + block_table, + sparse_count, +): + del key, weights, actual_seq_lengths_query, actual_seq_lengths_key + del block_table + return query.new_empty((query.shape[0], 1, sparse_count), dtype=torch.int32) + + +def _sparse_flash_attention_abstract( + query, + key, + value, + sparse_indices, + scale_value, + block_table, + actual_seq_lengths_query, + kv_seqlens, + query_rope, + key_rope, +): + del key, sparse_indices, scale_value, block_table + del actual_seq_lengths_query, kv_seqlens, query_rope, key_rope + return query.new_empty((*query.shape[:-1], value.shape[-1])) + + @register_custom_op("dlinfer::add_rms_norm", ["hidden_states", "residual"]) def add_rms_norm( hidden_states: Tensor, @@ -94,6 +130,107 @@ def apply_rotary_pos_emb( ) +@register_custom_op( + "dlinfer::apply_rotary_pos_emb_interleaved", + ["x"], + default_value={"return_native_layout": True}, +) +def apply_rotary_pos_emb_interleaved( + x: Tensor, + cos: Tensor, + sin: Tensor, + return_native_layout: bool, +) -> Tensor: + """Apply complex RoPE to adjacent element pairs. + + By default, the output keeps the native vendor front/back-half layout. + Set return_native_layout=False to restore the adjacent-pair layout. + The cos and sin tables use the front/back-half layout required by the + vendor implementation. + """ + impl = vendor_ops_registry.get("apply_rotary_pos_emb_interleaved") + if impl is not None: + return impl(x, cos, sin, return_native_layout) + + half_size = x.shape[-1] // 2 + x_even = x[..., 0::2] + x_odd = x[..., 1::2] + out_even = x_even * cos[..., :half_size] - x_odd * sin[..., :half_size] + out_odd = x_odd * cos[..., half_size:] + x_even * sin[..., half_size:] + if return_native_layout: + return torch.cat((out_even, out_odd), dim=-1) + return torch.stack((out_even, out_odd), dim=-1).flatten(-2) + + +@register_custom_op( + "dlinfer::lightning_indexer", + default_value={ + "actual_seq_lengths_query": None, + "actual_seq_lengths_key": None, + "block_table": None, + "sparse_count": 2048, + }, + impl_abstract_func=_lightning_indexer_abstract, +) +def lightning_indexer( + query: Tensor, + key: Tensor, + weights: Tensor, + actual_seq_lengths_query: Optional[Tensor], + actual_seq_lengths_key: Optional[Tensor], + block_table: Optional[Tensor], + sparse_count: int, +) -> Tensor: + """Select causal sparse-attention token indices from a paged key cache.""" + return vendor_ops_registry["lightning_indexer"]( + query, + key, + weights, + actual_seq_lengths_query, + actual_seq_lengths_key, + block_table, + sparse_count, + ) + + +@register_custom_op( + "dlinfer::sparse_flash_attention", + default_value={ + "block_table": None, + "actual_seq_lengths_query": None, + "kv_seqlens": None, + "query_rope": None, + "key_rope": None, + }, + impl_abstract_func=_sparse_flash_attention_abstract, +) +def sparse_flash_attention( + query: Tensor, + key: Tensor, + value: Tensor, + sparse_indices: Tensor, + scale_value: float, + block_table: Optional[Tensor], + actual_seq_lengths_query: Optional[Tensor], + kv_seqlens: Optional[Tensor], + query_rope: Optional[Tensor], + key_rope: Optional[Tensor], +) -> Tensor: + """Run sparse flash attention using logical token indices.""" + return vendor_ops_registry["sparse_flash_attention"]( + query, + key, + value, + sparse_indices, + scale_value, + block_table, + actual_seq_lengths_query, + kv_seqlens, + query_rope, + key_rope, + ) + + @register_custom_op( "dlinfer::prefill_attention", ["attn_output"], @@ -101,6 +238,7 @@ def apply_rotary_pos_emb( "softmax_scale": None, "alibi_slopes": None, "attn_output": None, + "actual_seq_lengths_cpu": None, }, ) def prefill_attention( @@ -109,7 +247,6 @@ def prefill_attention( value: Tensor, key_cache: Tensor, value_cache: Tensor, - q_start_loc: Tensor, q_seq_len: Tensor, kv_seq_len: Tensor, max_q_seq_len: int, @@ -119,6 +256,7 @@ def prefill_attention( softmax_scale: Optional[float], alibi_slopes: Optional[Sequence[float]], attn_output: Optional[Tensor], + actual_seq_lengths_cpu: Optional[Tensor], ) -> Tensor: """ Computes the multi-head attention over the query, key, and value tensors. @@ -130,7 +268,6 @@ def prefill_attention( value (Tensor): The value tensor. key_cache (Tensor): The existing key cache tensor. value_cache (Tensor): The existing value cache tensor. - q_start_loc (Tensor): The start location of each query sequence. q_seq_len (Tensor): The length of each query sequence. kv_seq_len (Tensor): The length of each key/value sequence. max_q_seq_len (int): The maximum length of any query sequence. @@ -140,6 +277,9 @@ def prefill_attention( softmax_scale (Optional[float]): The scale factor to apply to the attention logits before the softmax. alibi_slopes (Optional[Sequence[float]]): The slopes for the ALiBi attention bias, one for each head. attn_output (Optional[Tensor]): The computed attention output tensor. + actual_seq_lengths_cpu (Optional[Tensor]): Precomputed cumulative end + position of each query sequence on CPU. It is required for TND + attention. Returns: Tensor: The computed attention output tensor, alias of attn_output. @@ -148,7 +288,6 @@ def prefill_attention( query, key, value, - q_start_loc, q_seq_len, max_q_seq_len, num_q_heads, @@ -157,6 +296,7 @@ def prefill_attention( softmax_scale, alibi_slopes, attn_output, + actual_seq_lengths_cpu, ) @@ -348,10 +488,8 @@ def paged_prefill_attention( value_cache: Tensor, block_table: Tensor, block_size: int, - q_start_loc: Tensor, q_seq_len: Tensor, kv_seq_len: Tensor, - cu_seq_lens_kv: Tensor, max_q_seq_len: int, max_kv_seq_len: int, num_q_heads: int, @@ -376,10 +514,8 @@ def paged_prefill_attention( block_table (Tensor): A tensor that maps each position in the query sequence to the corresponding block in the key/value cache. block_size (int): The size of each block in the input sequence. - q_start_loc (Tensor): The start location of each query sequence. q_seq_len (Tensor): The length of each query sequence. kv_seq_len (Tensor): The length of each key/value sequence. - cu_seq_lens_kv (Tensor): The cumulative sequence lengths of the key/value sequences. max_q_seq_len (int): The maximum length of any query sequence. max_kv_seq_len (int): The maximum length of any key/value sequence. num_q_heads (int): The number of query heads. @@ -404,10 +540,8 @@ def paged_prefill_attention( value_cache, block_table, block_size, - q_start_loc, q_seq_len, kv_seq_len, - cu_seq_lens_kv, max_q_seq_len, max_kv_seq_len, num_q_heads, @@ -660,6 +794,39 @@ def fused_moe( return fused_moe_impl(*args, chunked_moe_layout) +def fused_moe_w8a8( + hidden_states: Tensor, + gate_up_weights: Tensor, + gate_up_scales: Tensor, + down_weights: Tensor, + down_scales: Tensor, + topk_weights: Tensor, + topk_ids: Tensor, + topk: int, + renormalize: bool, + moe_metadata: MoeMetadata, +) -> Tensor: + """Run dynamic W8A8 MoE through the vendor's non-fused fallback. + + Unlike :func:`fused_moe`, this entry point carries the per-channel weight + scales required by quantized grouped matmuls. It deliberately has no + ``chunked_moe_layout`` argument: GLM-5.2 has fewer than the Ascend grouped + matmul limit of 1024 (local) experts. + """ + return vendor_ops_registry["fused_moe_w8a8"]( + hidden_states, + gate_up_weights, + gate_up_scales, + down_weights, + down_scales, + topk_weights, + topk_ids, + topk, + renormalize, + moe_metadata, + ) + + def linear_impl_abstract_func( x: Tensor, weight: Tensor, @@ -788,6 +955,47 @@ def linear_w8a8( ) +def linear_w8a8_static_impl_abstract_func( + x: Tensor, + weight: Tensor, + input_scale: Tensor, + input_offset: Tensor, + deq_scale: Tensor, + out_dtype: torch.dtype, + quant_dtype: torch.dtype, + quant_bias: Optional[Tensor], +) -> Tensor: + return x.new_empty((*x.shape[:-1], weight.shape[0]), dtype=out_dtype) + + +@register_custom_op( + "dlinfer::linear_w8a8_static", + impl_abstract_func=linear_w8a8_static_impl_abstract_func, + default_value={"quant_bias": None}, +) +def linear_w8a8_static( + x: Tensor, + weight: Tensor, + input_scale: Tensor, + input_offset: Tensor, + deq_scale: Tensor, + out_dtype: torch.dtype, + quant_dtype: torch.dtype, + quant_bias: Optional[Tensor], +) -> Tensor: + """Run ModelSlim static W8A8 quantization followed by quant matmul.""" + return vendor_ops_registry["linear_w8a8_static"]( + x, + weight, + input_scale, + input_offset, + deq_scale, + out_dtype, + quant_dtype, + quant_bias, + ) + + def rms_norm_w8a8_impl_abstract_func( hidden_states: Tensor, weight: Tensor, diff --git a/dlinfer/vendor/ascend/attention.py b/dlinfer/vendor/ascend/attention.py index dd342448..e2781b23 100644 --- a/dlinfer/vendor/ascend/attention.py +++ b/dlinfer/vendor/ascend/attention.py @@ -13,8 +13,8 @@ def decode_attention( scale_value: float, block_table: Tensor, block_size: int, - q_seq_len: Tensor, - kv_seq_len: Tensor, + actual_q_seqlens_cpu: Tensor, + kv_seqlens_cpu: Tensor, softmax_scale: float, attn_output: Tensor, ): @@ -34,8 +34,8 @@ def decode_attention( block_table=block_table, input_layout="TND", block_size=block_size, - actual_seq_lengths=q_seq_len, - actual_seq_lengths_kv=kv_seq_len, + actual_seq_lengths=actual_q_seqlens_cpu, + actual_seq_lengths_kv=kv_seqlens_cpu, num_key_value_heads=num_kv_heads, num_heads=num_q_heads, scale=scale_value, @@ -51,7 +51,7 @@ def decode_attention_mla( num_q_heads: int, scale_value: float, block_table: Tensor, - kv_seq_len: Tensor, + kv_seqlens_cpu: Tensor, mla_vheadsize: int, attn_output: Tensor, ): @@ -85,7 +85,7 @@ def decode_attention_mla( block_table=block_table, block_size=block_size, actual_seq_qlen=None, - actual_seq_kvlen=kv_seq_len, + actual_seq_kvlen=kv_seqlens_cpu, ) attn_output.copy_(fai_output.squeeze(2).transpose(0, 1)) diff --git a/dlinfer/vendor/ascend/moe.py b/dlinfer/vendor/ascend/moe.py index a0a66033..a0486eac 100644 --- a/dlinfer/vendor/ascend/moe.py +++ b/dlinfer/vendor/ascend/moe.py @@ -107,6 +107,79 @@ def _grouped_mlp( return down_proj +def _grouped_mlp_w8a8( + hidden_states: torch.Tensor, + gate_up_weights: torch.Tensor, + gate_up_scales: torch.Tensor, + down_weights: torch.Tensor, + down_scales: torch.Tensor, + group_list: torch.Tensor, + group_list_type: int, +): + """Public-op W8A8 MoE fallback: quant GMM1, SwiGLU, quant GMM2. + + This intentionally uses only public NPU operators. The intermediate GMM1 + output is dequantized to the model dtype, then requantized per token after + SwiGLU, matching vllm-ascend's non-fused fallback. + """ + quantized, input_scale = torch.ops.npu.npu_dynamic_quant( + hidden_states, dst_type=torch.int8 + ) + gate_up = torch.ops.npu.npu_grouped_matmul( + [quantized], + [gate_up_weights.transpose(1, 2)], + scale=[gate_up_scales.squeeze(-1)], + per_token_scale=[input_scale], + group_list=group_list, + split_item=2, + group_type=0, + group_list_type=group_list_type, + output_dtype=hidden_states.dtype, + )[0] + + activated = torch.ops.npu.npu_swiglu(gate_up, -1) + activated_quant, activated_scale = torch.ops.npu.npu_dynamic_quant( + activated, dst_type=torch.int8 + ) + return torch.ops.npu.npu_grouped_matmul( + [activated_quant], + [down_weights.transpose(1, 2)], + scale=[down_scales.squeeze(-1)], + per_token_scale=[activated_scale], + group_list=group_list, + split_item=2, + group_type=0, + group_list_type=group_list_type, + output_dtype=hidden_states.dtype, + )[0] + + +def apply_mlp_w8a8( + hidden_states: torch.Tensor, + gate_up_weights: torch.Tensor, + gate_up_scales: torch.Tensor, + down_weights: torch.Tensor, + down_scales: torch.Tensor, + group_list: torch.Tensor, + group_list_type: int, +): + """Apply GLM-5.2 W8A8 experts through the non-fused public-op path.""" + if gate_up_weights.size(0) > MAX_GROUP_LIST_SIZE: + raise RuntimeError( + "W8A8 fallback supports at most 1024 local experts; " + f"got {gate_up_weights.size(0)}" + ) + return _grouped_mlp_w8a8( + hidden_states, + gate_up_weights, + gate_up_scales, + down_weights, + down_scales, + group_list, + group_list_type, + ) + + def _apply_mlp_chunked_eager( hidden_states: torch.Tensor, gate_up_weights: torch.Tensor, @@ -359,6 +432,8 @@ def fused_moe_naive( topk: int, renormalize: bool, chunked_moe_layout: ChunkedMoeWeightLayout = None, + gate_up_scales: torch.Tensor = None, + down_scales: torch.Tensor = None, ): num_experts = ( chunked_moe_layout.num_experts @@ -389,14 +464,25 @@ def fused_moe_naive( # MLP group_list_type = 1 expert_tokens = expert_tokens.to(torch.int64) - mlp_output = apply_mlp( - expanded_hidden_states, - gate_up_weights, - down_weights, - expert_tokens, - group_list_type, - chunked_moe_layout, - ) + if gate_up_scales is None: + mlp_output = apply_mlp( + expanded_hidden_states, + gate_up_weights, + down_weights, + expert_tokens, + group_list_type, + chunked_moe_layout, + ) + else: + mlp_output = apply_mlp_w8a8( + expanded_hidden_states, + gate_up_weights, + gate_up_scales, + down_weights, + down_scales, + expert_tokens, + group_list_type, + ) # distribute combine moe_output = torch.ops.npu.npu_moe_token_unpermute( @@ -419,6 +505,8 @@ def fused_moe_mc2( moe_group_name: str, x_active_mask: torch.Tensor, chunked_moe_layout: ChunkedMoeWeightLayout = None, + gate_up_scales: torch.Tensor = None, + down_scales: torch.Tensor = None, ): # do renormalize if renormalize: @@ -478,14 +566,25 @@ def fused_moe_mc2( # MLP group_list_type = 0 - mlp_output = apply_mlp( - expanded_hidden_states, - gate_up_weights, - down_weights, - expert_tokens, - group_list_type, - chunked_moe_layout, - ) + if gate_up_scales is None: + mlp_output = apply_mlp( + expanded_hidden_states, + gate_up_weights, + down_weights, + expert_tokens, + group_list_type, + chunked_moe_layout, + ) + else: + mlp_output = apply_mlp_w8a8( + expanded_hidden_states, + gate_up_weights, + gate_up_scales, + down_weights, + down_scales, + expert_tokens, + group_list_type, + ) # distribute combine kwargs_mc2 = { @@ -535,6 +634,8 @@ def fused_moe_all2all( ep_group: dist.ProcessGroup, expert_ids_per_ep_rank: torch.Tensor, chunked_moe_layout: ChunkedMoeWeightLayout = None, + gate_up_scales: torch.Tensor = None, + down_scales: torch.Tensor = None, ): num_local_experts = ( chunked_moe_layout.num_experts @@ -692,14 +793,26 @@ def fused_moe_all2all_forward( dispatched_outputs = dispatch(hidden_states, topk_ids) # MLP - mlp_output = apply_mlp( - dispatched_outputs["hidden_states"], - gate_up_weights, - down_weights, - dispatched_outputs["group_list"].to(torch.int64), - dispatched_outputs["group_list_type"], - chunked_moe_layout, - ) + group_list = dispatched_outputs["group_list"].to(torch.int64) + if gate_up_scales is None: + mlp_output = apply_mlp( + dispatched_outputs["hidden_states"], + gate_up_weights, + down_weights, + group_list, + dispatched_outputs["group_list_type"], + chunked_moe_layout, + ) + else: + mlp_output = apply_mlp_w8a8( + dispatched_outputs["hidden_states"], + gate_up_weights, + gate_up_scales, + down_weights, + down_scales, + group_list, + dispatched_outputs["group_list_type"], + ) # distribute combine context_metadata = dispatched_outputs["context_metadata"] diff --git a/dlinfer/vendor/ascend/torch_npu_ops.py b/dlinfer/vendor/ascend/torch_npu_ops.py index fa83b8c0..6d936ac5 100644 --- a/dlinfer/vendor/ascend/torch_npu_ops.py +++ b/dlinfer/vendor/ascend/torch_npu_ops.py @@ -15,13 +15,16 @@ MoECommType, MoeMetadata, ) -from .utils import SocVersion, get_cpu_seq_len +from .utils import SocVersion from .attention import decode_attention, decode_attention_mla from . import moe __all__ = [ "add_rms_norm", "apply_rotary_pos_emb", + "apply_rotary_pos_emb_interleaved", + "lightning_indexer", + "sparse_flash_attention", "prefill_attention", "incre_flash_attention", "fill_kv_cache", @@ -32,11 +35,13 @@ "get_cache_len", "weight_quant_matmul", "fused_moe", + "fused_moe_w8a8", "linear", "rms_norm_w8a8", "add_rms_norm_w8a8", "dynamic_quant", "linear_w8a8", + "linear_w8a8_static", ] @@ -89,9 +94,10 @@ def linear_w8a8( ) -> Tensor: out_dtype = torch.bfloat16 if out_dtype == torch.float16 else out_dtype - hidden_states = hidden_states.squeeze(0) - linear_scale = linear_scale.squeeze() - rms_scale = rms_scale.squeeze(0) + output_shape = (*hidden_states.shape[:-1], weight.shape[0]) + hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) + linear_scale = linear_scale.reshape(-1) + rms_scale = rms_scale.reshape(-1) output = torch.ops.npu.npu_quant_matmul( hidden_states, @@ -101,8 +107,40 @@ def linear_w8a8( bias=bias, output_dtype=out_dtype, ) - output = output.unsqueeze(0) - return output + return output.reshape(output_shape) + + +@register_ops(vendor_ops_registry) +def linear_w8a8_static( + hidden_states: Tensor, + weight: Tensor, + input_scale: Tensor, + input_offset: Tensor, + deq_scale: Tensor, + out_dtype: torch.dtype, + quant_dtype: torch.dtype, + quant_bias: Optional[Tensor], +) -> Tensor: + """ModelSlim static per-tensor activation / per-channel weight W8A8.""" + out_dtype = torch.bfloat16 if out_dtype == torch.float16 else out_dtype + input_size = hidden_states.shape[-1] + reciprocal_scale = input_scale.reciprocal().reshape(-1).expand(input_size) + offset = input_offset.to(input_scale.dtype).reshape(-1).expand(input_size) + quantized = torch.ops.npu.npu_quantize( + hidden_states, + reciprocal_scale, + offset, + torch.qint8, + -1, + False, + ) + return torch.ops.npu.npu_quant_matmul( + quantized, + weight.t(), + deq_scale.reshape(-1), + bias=quant_bias, + output_dtype=out_dtype, + ) @register_ops(vendor_ops_registry) @@ -168,12 +206,103 @@ def apply_rotary_pos_emb_(q, k, cos, sin): ) +@register_ops(vendor_ops_registry) +def apply_rotary_pos_emb_interleaved( + x: Tensor, + cos: Tensor, + sin: Tensor, + return_native_layout: bool = True, +) -> Tensor: + """Apply adjacent-pair RoPE with the native Ascend interleave operator. + + Return the vendor's native front/back-half layout by default. Set + return_native_layout=False to restore adjacent pairs. + """ + if x.ndim != 4 or cos.ndim != 4 or sin.ndim != 4: + raise ValueError( + "npu_interleave_rope expects 4D x/cos/sin tensors, got " + f"{x.ndim}D, {cos.ndim}D and {sin.ndim}D" + ) + output = torch_npu.npu_interleave_rope( + x.contiguous(), cos.contiguous(), sin.contiguous() + ) + if return_native_layout: + return output + + # The native op returns rotated even and odd elements in separate halves. + output_even, output_odd = output.chunk(2, dim=-1) + return torch.stack((output_even, output_odd), dim=-1).flatten(-2) + +@register_ops(vendor_ops_registry) +def lightning_indexer( + query: Tensor, + key: Tensor, + weights: Tensor, + actual_seq_lengths_query: Optional[Tensor], + actual_seq_lengths_key: Optional[Tensor], + block_table: Optional[Tensor], + sparse_count: int, +) -> Tensor: + """BF16 Lightning Indexer backed by the native torch-npu operator.""" + indices, _ = torch_npu.npu_lightning_indexer( + query=query.contiguous(), + key=key, + weights=weights.contiguous(), + actual_seq_lengths_query=actual_seq_lengths_query, + actual_seq_lengths_key=actual_seq_lengths_key, + block_table=block_table, + layout_query="TND", + layout_key="PA_BSND", + sparse_count=sparse_count, + sparse_mode=3, + ) + return indices + + +@register_ops(vendor_ops_registry) +def sparse_flash_attention( + query: Tensor, + key: Tensor, + value: Tensor, + sparse_indices: Tensor, + scale_value: float, + block_table: Optional[Tensor], + actual_seq_lengths_query: Optional[Tensor], + kv_seqlens: Optional[Tensor], + query_rope: Optional[Tensor], + key_rope: Optional[Tensor], +) -> Tensor: + """BF16 sparse flash attention backed by the native torch-npu operator.""" + if query.dtype not in (torch.bfloat16, torch.float16): + raise TypeError( + f"sparse_flash_attention expects BF16/FP16 query, got {query.dtype}" + ) + + output, _, _ = torch_npu.npu_sparse_flash_attention( + query=query.contiguous(), + key=key, + value=value, + sparse_indices=sparse_indices.contiguous(), + scale_value=scale_value, + block_table=block_table, + actual_seq_lengths_query=actual_seq_lengths_query, + actual_seq_lengths_kv=kv_seqlens, + query_rope=None if query_rope is None else query_rope.contiguous(), + key_rope=key_rope, + sparse_block_size=1, + layout_query="TND", + layout_kv="PA_BSND", + sparse_mode=3, + attention_mode=2, + ) + return output + + @register_ops(vendor_ops_registry) def prefill_attention( query: Tensor, key: Tensor, value: Tensor, - q_start_loc: Tensor, q_seq_len: Tensor, max_q_seq_len: int, num_q_heads: int, @@ -182,6 +311,7 @@ def prefill_attention( softmax_scale: Optional[float], alibi_slopes: Optional[Sequence[float]], attn_output: Optional[Tensor], + actual_seq_lengths_cpu: Optional[Tensor], ) -> Tensor: if alibi_slopes is not None: raise RuntimeError( @@ -197,14 +327,18 @@ def prefill_attention( mask = attn_mask[0] else: # Handle qwenvl vision part flash-attention - q_seq_len = get_cpu_seq_len(q_seq_len) is_tnd = ( query.dim() == 3 and query.shape[-2] == num_q_heads and key.shape[-2] == num_kv_heads ) input_layout = "TND" if is_tnd else "BSH" - actual_seq_lengths = q_seq_len.cumsum(dim=0) if is_tnd else None + if is_tnd and actual_seq_lengths_cpu is None: + raise ValueError( + "actual_seq_lengths_cpu is required for TND prefill attention" + ) + if not is_tnd: + actual_seq_lengths_cpu = None fia_kwargs = {} if is_tnd and query.shape[-1] > value.shape[-1]: nope_dim = value.shape[-1] @@ -217,8 +351,8 @@ def prefill_attention( key=key, value=value, input_layout=input_layout, - actual_seq_lengths=actual_seq_lengths, - actual_seq_lengths_kv=actual_seq_lengths, + actual_seq_lengths=actual_seq_lengths_cpu, + actual_seq_lengths_kv=actual_seq_lengths_cpu, scale=scale_value, num_heads=num_q_heads, num_key_value_heads=num_kv_heads, @@ -228,9 +362,6 @@ def prefill_attention( attn_output.copy_(output) return attn_output if SocVersion.is_Ascend910(): - q_seq_len = get_cpu_seq_len(q_seq_len) - actual_seq_lengths = q_seq_len.cumsum(dim=0) - # The backend supplies the fixed split-fuse causal mask required by # sparse mode 3 for both standard attention and MLA. fia_kwargs = {} @@ -249,8 +380,8 @@ def prefill_attention( value=value, atten_mask=mask, input_layout="TND", - actual_seq_lengths=actual_seq_lengths, - actual_seq_lengths_kv=actual_seq_lengths, + actual_seq_lengths=actual_seq_lengths_cpu, + actual_seq_lengths_kv=actual_seq_lengths_cpu, scale=scale_value, num_heads=num_q_heads, num_key_value_heads=num_kv_heads, @@ -343,6 +474,32 @@ def fill_kv_cache( v_scales_zeros: Sequence[Optional[Tensor]], quant_bits: int, ) -> Tuple[Tensor, Tensor]: + split_mla_cache = ( + key_cache.dim() == 4 + and value_cache.dim() == 4 + and key_cache.size(-1) > 0 + and value_cache.size(-1) > 0 + and key.size(-1) == key_cache.size(-1) + value_cache.size(-1) + and value.size(-1) == value_cache.size(-1) + ) + if split_mla_cache: + # key_cache -> rope_cache, value_cache -> nope_cache + rope_head_size = key_cache.size(-1) + key_rope = key[..., -rope_head_size:].contiguous() + value_nope = value.contiguous() + key_cache_reshaped = torch.flatten(key_cache, start_dim=0, end_dim=1) + value_cache_reshaped = torch.flatten( + value_cache, start_dim=0, end_dim=1 + ) + kv_indices = kv_indices.view(-1, 1) + torch.ops.npu.npu_scatter_nd_update_( + key_cache_reshaped, kv_indices, key_rope + ) + torch.ops.npu.npu_scatter_nd_update_( + value_cache_reshaped, kv_indices, value_nope + ) + return key_cache, value_cache + # only support contiguous k,v key = key.contiguous() value = value.contiguous() @@ -396,8 +553,8 @@ def paged_decode_attention( value_cache: Tensor, block_table: Optional[Tensor], block_size: int, - q_seq_len: Tensor, - kv_seq_len: Tensor, + actual_q_seqlens_cpu: Tensor, + kv_seqlens_cpu: Tensor, max_kv_seq_len: int, num_q_heads: int, num_kv_heads: int, @@ -429,8 +586,8 @@ def paged_decode_attention( scale_value=scale_value, block_table=block_table, block_size=block_size, - q_seq_len=q_seq_len, - kv_seq_len=kv_seq_len, + actual_q_seqlens_cpu=actual_q_seqlens_cpu, + kv_seqlens_cpu=kv_seqlens_cpu, softmax_scale=softmax_scale, attn_output=attn_output, ) @@ -442,7 +599,7 @@ def paged_decode_attention( num_q_heads=num_q_heads, scale_value=scale_value, block_table=block_table, - kv_seq_len=kv_seq_len, + kv_seqlens_cpu=kv_seqlens_cpu, mla_vheadsize=value_headsize, attn_output=attn_output, ) @@ -457,10 +614,8 @@ def paged_prefill_attention( value_cache: Tensor, block_table: Tensor, block_size: int, - q_start_loc: Tensor, - q_seq_len: Tensor, - kv_seq_len: Tensor, - cu_seq_lens_kv: Tensor, + actual_q_seqlens_cpu: Tensor, + kv_seqlens_cpu: Tensor, max_q_seq_len: int, max_kv_seq_len: int, num_q_heads: int, @@ -525,8 +680,8 @@ def paged_prefill_attention( softmax_scale=scale_value, block_table=block_table, block_size=block_size, - actual_seq_qlen=q_seq_len, - actual_seq_kvlen=kv_seq_len, + actual_seq_qlen=actual_q_seqlens_cpu, + actual_seq_kvlen=kv_seqlens_cpu, ) # TND_NTD returns [num_heads, num_tokens, value_head_size]. @@ -548,8 +703,8 @@ def paged_prefill_attention( block_table=block_table, input_layout="TND", block_size=block_size, - actual_seq_lengths=q_seq_len, - actual_seq_lengths_kv=kv_seq_len, + actual_seq_lengths=actual_q_seqlens_cpu, + actual_seq_lengths_kv=kv_seqlens_cpu, num_key_value_heads=num_kv_heads, num_heads=num_q_heads, scale=scale_value, @@ -614,77 +769,6 @@ def moe_gating_topk_softmax( return routing_weights, selected_idx -# TODO only for internlm in transformers lib. -# see issue #9 for details -@register_ops(vendor_ops_registry) -def fused_attention( - query_states: Tensor, - key_states: Tensor, - value_states: Tensor, - mask: Sequence[Optional[Tensor]], -) -> Tensor: - batch_size = query_states.shape[0] - query_states = query_states.squeeze(0) - key_states = key_states.squeeze(0) - value_states = value_states.squeeze(0) - q_seq_len, num_q_heads, _ = query_states.shape - kv_seq_len, num_kv_heads, _ = value_states.shape - attn_output = torch.empty_like(query_states) - - for i in range(batch_size): - if q_seq_len == kv_seq_len: - # mask must be a square - if not mask[i : i + 1][0].shape[-1] == mask[i : i + 1][0].shape[-2]: - min_shape = min( - mask[i : i + 1][0].shape[-1], mask[i : i + 1][0].shape[-2] - ) - square_mask = mask[i : i + 1][0][..., :min_shape, :min_shape] - square_mask = square_mask.contiguous() - else: - square_mask = mask[i : i + 1][0] - - prefill_attention( - query_states, - key_states, - value_states, - torch.tensor( - [kv_seq_len - q_seq_len], - dtype=torch.int64, - device=query_states.device, - ), - torch.tensor( - [kv_seq_len], dtype=torch.int64, device=query_states.device - ), - q_seq_len, - num_q_heads, - num_kv_heads, - [ - square_mask, - ], - None, - None, - attn_output, - ) - else: - paged_decode_attention( - query_states, - key_states, - value_states, - None, - 0, - torch.tensor( - [kv_seq_len], dtype=torch.int64, device=query_states.device - ), - kv_seq_len, - num_q_heads, - num_kv_heads, - None, - None, - attn_output, - ) - return attn_output - - # Quantification of W4A16 is currently supported and tested. @register_ops(vendor_ops_registry) def weight_quant_matmul( @@ -792,6 +876,96 @@ def fused_moe( return moe_output +@register_ops(vendor_ops_registry) +def fused_moe_w8a8( + hidden_states: Tensor, + gate_up_weights: Tensor, + gate_up_scales: Tensor, + down_weights: Tensor, + down_scales: Tensor, + topk_weights: Tensor, + topk_ids: Tensor, + topk: int, + renormalize: bool, + moe_metadata: MoeMetadata, +) -> Tensor: + """Dynamic W8A8 MoE using only public torch_npu operators.""" + topk_ids = topk_ids.to(torch.int32) + ( + hidden_states, + num_tokens, + paded_num_tokens, + x_active_mask, + topk_ids, + topk_weights, + ) = moe.moe_prepare( + hidden_states, + moe_metadata.x_active_mask, + moe_metadata.pad_size, + moe_metadata.tp_size, + moe_metadata.ep_size, + moe_metadata.tp_rank, + moe_metadata.moe_comm_type, + topk_ids, + topk_weights, + ) + + if moe_metadata.moe_comm_type == MoECommType.MC2: + moe_output = moe.fused_moe_mc2( + hidden_states, + gate_up_weights, + down_weights, + topk_weights, + topk_ids, + renormalize, + moe_metadata.ep_size, + moe_metadata.ep_rank, + moe_metadata.moe_group_name, + x_active_mask, + None, + gate_up_scales, + down_scales, + ) + elif moe_metadata.moe_comm_type == MoECommType.ALLTOALL: + moe_output = moe.fused_moe_all2all( + hidden_states, + gate_up_weights, + down_weights, + topk_weights, + topk_ids, + renormalize, + moe_metadata.ep_size, + moe_metadata.ep_rank, + moe_metadata.ep_group, + moe_metadata.expert_ids_per_ep_rank, + None, + gate_up_scales, + down_scales, + ) + else: + moe_output = moe.fused_moe_naive( + hidden_states, + gate_up_weights, + down_weights, + topk_weights, + topk_ids, + topk, + renormalize, + None, + gate_up_scales, + down_scales, + ) + + return moe.moe_finalize( + moe_output, + num_tokens, + paded_num_tokens, + moe_metadata.ep_size, + moe_metadata.tp_size, + moe_metadata.tp_group, + ) + + @register_ops(vendor_ops_registry) def linear( x: Tensor, diff --git a/tests/test_ascend_attention_precision.py b/tests/test_ascend_attention_precision.py index abb43f71..9f85f92d 100644 --- a/tests/test_ascend_attention_precision.py +++ b/tests/test_ascend_attention_precision.py @@ -143,6 +143,7 @@ def _assert_prefill_attention_matches_torch( key = key.to(DEVICE) value = value.to(DEVICE) seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.int32) + actual_seq_lengths_cpu = seq_lens_tensor.cumsum(dim=0, dtype=torch.int32) max_seq_len = max(seq_lens) output = torch.empty( (num_tokens, NUM_Q_HEADS, value_head_dim), dtype=DTYPE, device=DEVICE @@ -152,7 +153,6 @@ def _assert_prefill_attention_matches_torch( query=query, key=key, value=value, - q_start_loc=None, q_seq_len=seq_lens_tensor, max_q_seq_len=max_seq_len, num_q_heads=NUM_Q_HEADS, @@ -161,6 +161,7 @@ def _assert_prefill_attention_matches_torch( softmax_scale=softmax_scale, alibi_slopes=None, attn_output=output, + actual_seq_lengths_cpu=actual_seq_lengths_cpu, ) assert actual.data_ptr() == output.data_ptr() @@ -229,10 +230,8 @@ def test_paged_prefill_attention_mla_matches_torch(fai_causal_mask): value_cache=value_cache, block_table=block_table.to(DEVICE), block_size=block_size, - q_start_loc=None, q_seq_len=torch.tensor(cumulative_q_seq_lens, dtype=torch.int32), kv_seq_len=torch.tensor(kv_seq_lens, dtype=torch.int32), - cu_seq_lens_kv=None, max_q_seq_len=max(q_seq_lens), max_kv_seq_len=max(kv_seq_lens), num_q_heads=NUM_Q_HEADS, @@ -279,10 +278,8 @@ def test_paged_prefill_attention_mla_graph_replay(fai_causal_mask): value_cache=key_cache[..., :MLA_V_HEAD_DIM], block_table=block_table.to(DEVICE), block_size=block_size, - q_start_loc=None, q_seq_len=torch.tensor(q_seq_lens, dtype=torch.int32), kv_seq_len=torch.tensor(capture_kv_seq_lens, dtype=torch.int32), - cu_seq_lens_kv=None, max_q_seq_len=max(q_seq_lens), max_kv_seq_len=max(capture_kv_seq_lens), num_q_heads=NUM_Q_HEADS, @@ -349,7 +346,7 @@ def test_decode_attention_mla_matches_torch(): num_q_heads=NUM_Q_HEADS, scale_value=MLA_SOFTMAX_SCALE, block_table=block_table.to(DEVICE), - kv_seq_len=torch.tensor(kv_seq_lens, dtype=torch.int32), + kv_seqlens_cpu=torch.tensor(kv_seq_lens, dtype=torch.int32), mla_vheadsize=MLA_V_HEAD_DIM, attn_output=output, ) diff --git a/tests/test_ascend_cudagraph_nsa.py b/tests/test_ascend_cudagraph_nsa.py new file mode 100644 index 00000000..a3afced7 --- /dev/null +++ b/tests/test_ascend_cudagraph_nsa.py @@ -0,0 +1,133 @@ +# Copyright (c) 2026, DeepLink. All rights reserved. + +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("torch_npu") + +from dlinfer.framework.lmdeploy_ext.cudagraph import ascend_cudagraph + + +def _graph_meta(max_batches=4, query_len=1, sparse=True): + return SimpleNamespace( + max_batchs=max_batches, + max_tokens=max_batches * query_len, + num_blocks=3, + device=torch.device("cpu"), + is_ssm=False, + use_mrope=False, + vocab_size=128, + mla_index_topk=2048 if sparse else None, + input_buffers={}, + ) + + +def _model(model_type="glm_moe_dsa"): + return SimpleNamespace(config=SimpleNamespace(model_type=model_type)) + + +@pytest.mark.parametrize("query_len", [1, 3]) +def test_graph_buffers_use_canonical_sequence_metadata(query_len): + graph_meta = _graph_meta(query_len=query_len) + buffers = ascend_cudagraph.AscendCudaGraphMixin_make_buffers_cudagraph( + _model(), graph_meta + ) + + assert buffers["q_seqlens"].tolist() == [query_len] * 4 + assert buffers["cu_seqlens_q"].tolist() == [ + step * query_len for step in range(5) + ] + assert buffers["kv_seqlens"].tolist() == [1, 1, 1, 1] + assert "cu_seqlens_q_cpu" not in buffers + assert "kv_seqlens_cpu" not in buffers + assert "attention_mask" not in buffers + + +def test_non_dsa_graph_uses_same_metadata_contract(): + buffers = ascend_cudagraph.AscendCudaGraphMixin_make_buffers_cudagraph( + _model("deepseek_v2"), _graph_meta(sparse=False) + ) + + assert not any(name.startswith("nsa_") for name in buffers) + assert set(("q_seqlens", "cu_seqlens_q", "cu_seqlens_q_cpu", + "kv_seqlens", "kv_seqlens_cpu")) <= buffers.keys() + assert "attention_mask" in buffers + + +def test_fill_glm_dsa_graph_buffers_pads_and_rebinds_metadata(monkeypatch): + graph_meta = _graph_meta() + graph_meta.input_buffers = ( + ascend_cudagraph.AscendCudaGraphMixin_make_buffers_cudagraph( + _model(), graph_meta + ) + ) + moe_metadata = SimpleNamespace(x_active_mask=None) + context = SimpleNamespace(moe_metadata=moe_metadata) + manager = SimpleNamespace(current_context=lambda: context) + monkeypatch.setattr(ascend_cudagraph, "get_step_ctx_manager", lambda: manager) + + metadata = SimpleNamespace( + block_offsets=torch.tensor([[7, 8], [9, 10]], dtype=torch.int32), + kv_seqlens=torch.tensor([5, 7], dtype=torch.int32), + kv_seqlens_cpu=None, + kv_start_indices=torch.tensor([4, 6], dtype=torch.int32), + q_start_loc=None, + cache_seqlens=None, + is_multi_token_decoding=False, + cu_seqlens_q=torch.tensor([0, 1, 2], dtype=torch.int32), + cu_seqlens_q_cpu=None, + ) + inputs = ascend_cudagraph.AscendCudaGraphMixin_fill_buffers_cudagraph( + _model(), + graph_meta, + input_ids=torch.tensor([[11, 12]], dtype=torch.int32), + position_ids=torch.tensor([[4, 6]], dtype=torch.int32), + past_key_values=[], + attn_metadata=metadata, + inputs_embeds=None, + ) + + assert metadata.q_seqlens is graph_meta.input_buffers["q_seqlens"] + assert metadata.cu_seqlens_q is graph_meta.input_buffers["cu_seqlens_q"] + assert metadata.cu_seqlens_q_cpu is None + assert metadata.kv_seqlens is graph_meta.input_buffers["kv_seqlens"] + assert metadata.kv_seqlens_cpu is None + assert metadata.q_seqlens.tolist() == [1, 1, 1, 1] + assert metadata.cu_seqlens_q.tolist() == [0, 1, 2, 3, 4] + assert metadata.kv_seqlens.tolist() == [5, 7, 0, 0] + assert inputs["attn_metadata"] is metadata + + +def test_sparse_graph_replay_skips_cpu_metadata_update(): + class _Graph: + replayed = False + + def replay(self): + self.replayed = True + + def update(self, **kwargs): + pytest.fail("sparse graph must not update CPU attention metadata") + + output = object() + graph = _Graph() + runner = object.__new__(ascend_cudagraph.AscendSingleGraphRunner) + runner._graph = graph + runner.meta = SimpleNamespace( + mla_index_topk=2048, + input_buffers={}, + output_buffers=output, + ) + runner.model = SimpleNamespace( + fill_buffers_cudagraph=lambda *args, **kwargs: None, + update_context_cudagraph=lambda *args, **kwargs: None, + get_outputs_cudagraph=lambda buffers, **kwargs: buffers, + ) + runner.ctx_mgr = SimpleNamespace(current_context=lambda: object()) + runner.is_mla = True + + actual = runner.forward() + + assert graph.replayed + assert actual is output diff --git a/tests/test_ascend_interleaved_rope.py b/tests/test_ascend_interleaved_rope.py new file mode 100644 index 00000000..98fa9222 --- /dev/null +++ b/tests/test_ascend_interleaved_rope.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026, DeepLink. All rights reserved. + +import pytest +import torch + +torch_npu = pytest.importorskip("torch_npu") + +if not torch.npu.is_available(): + pytest.skip("Ascend NPU is required", allow_module_level=True) + +from dlinfer.ops import apply_rotary_pos_emb_interleaved + + +def _reference(x, cos, sin): + x_even = x[..., 0::2] + x_odd = x[..., 1::2] + return torch.stack( + (x_even * cos - x_odd * sin, x_odd * cos + x_even * sin), dim=-1 + ).flatten(-2) + + +@pytest.mark.parametrize("num_heads", [1, 3, 32]) +def test_interleaved_rope_matches_adjacent_pair_reference(num_heads): + torch.manual_seed(20260821) + num_tokens = 7 + head_dim = 64 + x = torch.randn( + num_tokens, num_heads, 1, head_dim, dtype=torch.bfloat16, device="npu" + ) + cos = torch.randn( + num_tokens, 1, 1, head_dim // 2, dtype=torch.bfloat16, device="npu" + ) + sin = torch.randn_like(cos) + cos_native = torch.cat((cos, cos), dim=-1) + sin_native = torch.cat((sin, sin), dim=-1) + + expected = _reference(x, cos, sin) + actual_native = apply_rotary_pos_emb_interleaved(x, cos_native, sin_native) + actual_adjacent = apply_rotary_pos_emb_interleaved( + x, cos_native, sin_native, return_native_layout=False + ) + expected_native = torch.cat( + (expected[..., 0::2], expected[..., 1::2]), dim=-1 + ) + + for actual in (actual_native, actual_adjacent): + assert actual.shape == x.shape + assert actual.dtype == x.dtype + torch.testing.assert_close( + actual_native.float(), expected_native.float(), rtol=2e-2, atol=2e-2 + ) + torch.testing.assert_close( + actual_adjacent.float(), expected.float(), rtol=2e-2, atol=2e-2 + ) diff --git a/tests/test_ascend_sparse_attention.py b/tests/test_ascend_sparse_attention.py new file mode 100644 index 00000000..0ca9be49 --- /dev/null +++ b/tests/test_ascend_sparse_attention.py @@ -0,0 +1,202 @@ +# Copyright (c) 2026, DeepLink. All rights reserved. + +import pytest +import torch + +torch_npu = pytest.importorskip("torch_npu") + +if not torch.npu.is_available(): + pytest.skip("Ascend NPU is required", allow_module_level=True) + +from dlinfer.ops import fill_kv_cache, lightning_indexer, sparse_flash_attention + +DTYPE = torch.bfloat16 +DEVICE = torch.device("npu") +BLOCK_SIZE = 128 +TOPK = 2048 + + +def _randn(shape): + return torch.randn(shape, dtype=torch.float32).to(device=DEVICE, dtype=DTYPE) + + +def _metadata(): + cumulative_q = torch.tensor([3, 5], dtype=torch.int32) + kv_lengths = torch.tensor([3, 2], dtype=torch.int32) + block_table = torch.tensor([[0], [1]], dtype=torch.int32) + return cumulative_q, kv_lengths, block_table + + +def _sparse_attention_reference( + query, query_rope, combined_cache, indices, cumulative_q, kv_lengths, block_table +): + query = query.float().cpu() + query_rope = query_rope.float().cpu() + combined_cache = combined_cache.float().cpu() + indices = indices.cpu() + cumulative_q = cumulative_q.tolist() + kv_lengths = kv_lengths.tolist() + block_table = block_table.tolist() + output = torch.empty_like(query) + q_start = 0 + + for seq_idx, q_end in enumerate(cumulative_q): + kv_length = kv_lengths[seq_idx] + num_blocks = (kv_length + BLOCK_SIZE - 1) // BLOCK_SIZE + cache = torch.cat( + [combined_cache[block_table[seq_idx][block]] for block in range(num_blocks)] + )[:kv_length, 0] + key = cache[:, : query.size(-1)] + key_rope = cache[:, query.size(-1) :] + for token_idx in range(q_start, q_end): + selected = indices[token_idx, 0] + selected = selected[(selected >= 0) & (selected < kv_length)].long() + scores = torch.einsum("hd,kd->hk", query[token_idx], key[selected]) + scores += torch.einsum( + "hd,kd->hk", query_rope[token_idx], key_rope[selected] + ) + probs = torch.softmax(scores * (query.size(-1) ** -0.5), dim=-1) + output[token_idx] = torch.einsum("hk,kd->hd", probs, key[selected]) + q_start = q_end + return output + + +def test_lightning_indexer_matches_native_torch_npu(): + torch.manual_seed(20260817) + query = _randn((5, 32, 128)) + key_cache = _randn((2, BLOCK_SIZE, 1, 128)) + weights = _randn((5, 32)) + cumulative_q, kv_lengths, block_table = _metadata() + cumulative_q_device = cumulative_q.to(DEVICE) + kv_lengths_device = kv_lengths.to(DEVICE) + block_table_device = block_table.to(DEVICE) + + expected, _ = torch_npu.npu_lightning_indexer( + query=query, + key=key_cache, + weights=weights, + actual_seq_lengths_query=cumulative_q_device, + actual_seq_lengths_key=kv_lengths_device, + block_table=block_table_device, + layout_query="TND", + layout_key="PA_BSND", + sparse_count=TOPK, + sparse_mode=3, + ) + actual = lightning_indexer( + query, + key_cache, + weights, + actual_seq_lengths_query=cumulative_q_device, + actual_seq_lengths_key=kv_lengths_device, + block_table=block_table_device, + sparse_count=TOPK, + ) + + torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=0, atol=0) + assert actual.shape == (5, 1, TOPK) + assert actual.dtype == torch.int32 + + +def test_sparse_flash_attention_matches_native_torch_npu(): + torch.manual_seed(20260818) + cumulative_q, kv_lengths, block_table = _metadata() + cumulative_q_device = cumulative_q.to(DEVICE) + kv_lengths_device = kv_lengths.to(DEVICE) + block_table_device = block_table.to(DEVICE) + index_query = _randn((5, 32, 128)) + index_cache = _randn((2, BLOCK_SIZE, 1, 128)) + index_weights = _randn((5, 32)) + indices, _ = torch_npu.npu_lightning_indexer( + query=index_query, + key=index_cache, + weights=index_weights, + actual_seq_lengths_query=cumulative_q_device, + actual_seq_lengths_key=kv_lengths_device, + block_table=block_table_device, + layout_query="TND", + layout_key="PA_BSND", + sparse_count=TOPK, + sparse_mode=3, + ) + + query = _randn((5, 64, 512)) + query_rope = _randn((5, 64, 64)) + combined_cache = _randn((2, BLOCK_SIZE, 1, 576)) + key = combined_cache[..., :512] + key_rope = combined_cache[..., 512:] + + expected, _, _ = torch_npu.npu_sparse_flash_attention( + query=query, + key=key, + value=key, + sparse_indices=indices, + scale_value=512**-0.5, + block_table=block_table_device, + actual_seq_lengths_query=cumulative_q_device, + actual_seq_lengths_kv=kv_lengths_device, + query_rope=query_rope, + key_rope=key_rope, + sparse_block_size=1, + layout_query="TND", + layout_kv="PA_BSND", + sparse_mode=3, + attention_mode=2, + ) + actual = sparse_flash_attention( + query, + key, + key, + indices, + 512**-0.5, + block_table=block_table_device, + actual_seq_lengths_query=cumulative_q_device, + kv_seqlens=kv_lengths_device, + query_rope=query_rope, + key_rope=key_rope, + ) + + torch.testing.assert_close( + actual.cpu().float(), expected.cpu().float(), rtol=5e-3, atol=5e-3 + ) + reference = _sparse_attention_reference( + query, + query_rope, + combined_cache, + indices, + cumulative_q, + kv_lengths, + block_table, + ) + torch.testing.assert_close(actual.cpu().float(), reference, rtol=3e-2, atol=3e-2) + + +def test_fill_kv_cache_writes_split_mla_caches(): + torch.manual_seed(20260828) + key = _randn((3, 1, 576)) + value = _randn((3, 1, 512)) + rope_cache = torch.zeros( + (2, BLOCK_SIZE, 1, 64), dtype=DTYPE, device=DEVICE + ) + nope_cache = torch.zeros( + (2, BLOCK_SIZE, 1, 512), dtype=DTYPE, device=DEVICE + ) + slot_indices = torch.tensor( + [0, 7, BLOCK_SIZE + 3], dtype=torch.int32, device=DEVICE + ) + + actual_rope, actual_nope = fill_kv_cache( + key, + value, + rope_cache, + nope_cache, + slot_indices, + k_scales_zeros=(), + v_scales_zeros=(), + quant_bits=0, + ) + + flat_rope = actual_rope.flatten(0, 1) + flat_nope = actual_nope.flatten(0, 1) + torch.testing.assert_close(flat_rope[slot_indices].cpu(), key[..., -64:].cpu()) + torch.testing.assert_close(flat_nope[slot_indices].cpu(), value.cpu())