dlinfer-ascend 0.2.3.post2__cp311-cp311-manylinux2014_aarch64.whl

dlinfer/framework/lmdeploy_ext/cudagraph/ppu_cudagraph.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2025, OpenMMLab and DeepLink. All rights reserved.
+import torch
+from torch import Tensor
+
+from typing import Any, Dict, List
+
+from lmdeploy.pytorch.models.utils.cudagraph import CudaGraphMeta
+from lmdeploy.pytorch.models.utils.cudagraph import CudaGraphMixin, next_power_of_2
+
+BuffType = Dict[str, Tensor]
+
+
+def PPUCudaGraphMixin_make_buffers_cudagraph(
+    self, graph_meta: CudaGraphMeta, *args, **kwargs
+) -> BuffType:
+    """make cudagraph buffers from forward inputs."""
+    max_batches = graph_meta.max_batchs
+    max_tokens = graph_meta.max_tokens
+    num_blocks = graph_meta.num_blocks
+    device = graph_meta.device
+    input_buffers: BuffType = dict()
+    input_buffers["input_ids"] = torch.zeros(
+        1, max_tokens, dtype=torch.int32, device=device
+    )
+
+    input_buffers["position_ids"] = torch.zeros(
+        (1, max_tokens), dtype=torch.int32, device=device
+    )
+
+    input_buffers["block_offsets"] = torch.zeros(
+        (max_batches, num_blocks), dtype=torch.int32, device=device
+    )
+
+    input_buffers["q_seqlens"] = torch.zeros(
+        max_batches, dtype=torch.int32, device=device
+    )
+
+    input_buffers["kv_seqlens"] = torch.zeros(
+        max_batches, dtype=torch.int32, device=device
+    )
+
+    input_buffers["q_start_loc"] = torch.zeros(
+        max_batches + 1, dtype=torch.int32, device=device
+    )
+
+    input_buffers["kv_start_indices"] = torch.zeros(
+        (max_batches, 1), dtype=torch.int64, device=device
+    )
+    return input_buffers
+
+
+def PPUCudaGraphMixin_fill_buffers_cudagraph(
+    self,
+    graph_meta: CudaGraphMeta,
+    input_ids: Tensor,
+    position_ids: Tensor,
+    past_key_values: List,
+    attn_metadata: Any,
+    inputs_embeds: Tensor,
+    **kwargs
+) -> Dict[str, Tensor]:
+    """fill cudagraph buffers from forward inputs."""
+    block_offsets: Tensor = attn_metadata.block_offsets
+    q_start_loc: Tensor = attn_metadata.q_start_loc
+    q_seqlens: Tensor = attn_metadata.q_seqlens
+    kv_seqlens: Tensor = attn_metadata.kv_seqlens
+    kv_start_indices: Tensor = attn_metadata.kv_start_indices
+
+    input_buffers: BuffType = graph_meta.input_buffers
+
+    batch_size, num_blocks = block_offsets.size()
+    num_tokens = input_ids.size(-1)
+
+    # fill buffer
+    input_buffers["input_ids"][:, :num_tokens] = input_ids
+    input_buffers["position_ids"][:, :num_tokens] = position_ids
+    input_buffers["block_offsets"][:batch_size, :num_blocks] = block_offsets
+    input_buffers["q_start_loc"][: batch_size + 1] = q_start_loc
+    input_buffers["q_seqlens"][:batch_size] = q_seqlens
+    input_buffers["kv_seqlens"][:batch_size] = kv_seqlens
+    input_buffers["kv_start_indices"][:batch_size] = kv_start_indices
+
+    if inputs_embeds is not None:
+        emb_size = inputs_embeds.size(-1)
+        if "inputs_embeds" not in input_buffers:
+            max_num_tokens = input_buffers["input_ids"].size(-1)
+            input_buffers["inputs_embeds"] = inputs_embeds.new_zeros(
+                1, max_num_tokens, emb_size
+            )
+        input_buffers["inputs_embeds"][:, :num_tokens] = inputs_embeds
+
+    # create inputs
+    new_batch_size = next_power_of_2(batch_size)
+
+    attn_metadata.block_offsets = input_buffers["block_offsets"][:new_batch_size]
+    attn_metadata.q_start_loc = input_buffers["q_start_loc"][: new_batch_size + 1]
+    attn_metadata.q_seqlens = input_buffers["q_seqlens"][:new_batch_size]
+    attn_metadata.kv_seqlens = input_buffers["kv_seqlens"][:new_batch_size]
+    attn_metadata.kv_start_indices = input_buffers["kv_start_indices"][:new_batch_size]
+    attn_metadata.max_q_seq_len = int(attn_metadata.q_seqlens.max().item())
+    attn_metadata.max_kv_seq_len = int(attn_metadata.kv_seqlens.max().item())
+
+    new_inputs = dict(
+        past_key_values=past_key_values,
+        attn_metadata=attn_metadata,
+    )
+
+    new_inputs["input_ids"] = input_buffers["input_ids"][:, :new_batch_size]
+    new_inputs["position_ids"] = input_buffers["position_ids"][:, :new_batch_size]
+
+    if inputs_embeds is not None:
+        new_inputs["inputs_embeds"] = input_buffers["inputs_embeds"][:, :new_batch_size]
+
+    new_inputs.update(kwargs)
+
+    return new_inputs
+
+
+def PPUCudaGraphMixin_update_context_cudagraph(self, graph_meta, context):
+    """update step context with input buffers."""
+    input_buffers = graph_meta.input_buffers
+    context.block_offsets = input_buffers["block_offsets"]
+    context.q_seqlens = input_buffers["q_seqlens"]
+    context.kv_seqlens = input_buffers["kv_seqlens"]
+    context.q_start_loc = input_buffers["q_start_loc"]
+    context.kv_start_indices = input_buffers["kv_start_indices"]
+
+
+CudaGraphMixin.make_buffers_cudagraph = PPUCudaGraphMixin_make_buffers_cudagraph
+CudaGraphMixin.fill_buffers_cudagraph = PPUCudaGraphMixin_fill_buffers_cudagraph
+CudaGraphMixin.update_context_cudagraph = PPUCudaGraphMixin_update_context_cudagraph

dlinfer/framework/lmdeploy_ext/device/__init__.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2024, DeepLink. All rights reserved.
+import importlib
+import torch
+from functools import lru_cache
+from dlinfer.vendor import vendor_name
+from lmdeploy.pytorch import models
+
+
+vendor = ["camb", "ascend"]
+
+
+def fake_torch_compile(dynamic=False):
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def pre_rms_norm(q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
+    """Pre rms norm."""
+    q = q.to(torch.float32)
+    k = k.to(torch.float32)
+    variance_q = (q * q).sum(-1, keepdim=True)
+    variance_k = (k * k).sum(-1, keepdim=True)
+    variance = torch.stack([variance_q, variance_k], dim=0)
+    return variance
+
+
+def post_rms_norm(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    weight_q: torch.Tensor,
+    weight_k: torch.Tensor,
+    variance: torch.Tensor,
+    eps: float,
+    embed_dim: int,
+    dtype: torch.dtype,
+):
+    """Post rms norm."""
+    q = q.to(torch.float32)
+    k = k.to(torch.float32)
+    variance = variance / embed_dim + eps
+    variance_q, variance_k = variance
+    q = q * torch.rsqrt(variance_q)
+    q = q.to(dtype) * weight_q
+    k = k * torch.rsqrt(variance_k)
+    k = k.to(dtype) * weight_k
+    return q, k
+
+
+def patch_compiled_func():
+    import torch
+
+    real_torch_compile = torch.compile
+    torch.compile = fake_torch_compile
+    from lmdeploy.pytorch.models import internvl, internvl3_hf
+
+    internvl.pre_rms_norm = pre_rms_norm
+    internvl.post_rms_norm = post_rms_norm
+    internvl3_hf.pre_rms_norm = pre_rms_norm
+    internvl3_hf.post_rms_norm = post_rms_norm
+    torch.compile = real_torch_compile
+
+
+@lru_cache(1)
+def import_vendor_module(vendor_name_str):
+    if vendor_name_str in vendor:
+        importlib.import_module(f".{vendor_name_str}", __package__)
+
+
+def vendor_device_init():
+    import_vendor_module(vendor_name)
+    patch_compiled_func()
+
+
+vendor_device_init()

dlinfer/framework/lmdeploy_ext/device/ascend.py

@@ -0,0 +1,205 @@
+# Copyright (c) 2024, DeepLink. All rights reserved.
+import os
+import torch
+
+from lmdeploy.pytorch.backends.dlinfer.moe import DlinferFusedMoEImpl
+from lmdeploy.pytorch.models.chatglm2 import SelfAttention
+from lmdeploy.pytorch.engine import logits_process
+
+from dlinfer.vendor.ascend.utils import SocVersion
+
+
+def rl_update_weights(self, gate_up_weights: torch.Tensor, down_weights: torch.Tensor):
+    """Update weights."""
+    return gate_up_weights, down_weights
+
+
+if os.getenv("DLINFER_RESET_MOE_UPDATE_WEIGHTS", "0") == "1":
+    DlinferFusedMoEImpl.update_weights = rl_update_weights
+
+
+@staticmethod
+def ascend_chatglm2_fill_rope(states: torch.Tensor, rope: torch.Tensor):
+    """fill rope."""
+    rope_part = states.chunk(2, -1)[1]
+    rope = rope.unflatten(-1, (2, -1))
+    rope = rope.transpose(-2, -1).flatten(-2, -1)
+    states = torch.cat([rope_part, rope], dim=-1)
+
+    return states
+
+
+SelfAttention._fill_rope = ascend_chatglm2_fill_rope
+
+
+# modify bad words process for aclgraph performance
+def _process_bad_words_(
+    scores: torch.Tensor,
+    bad_words: torch.LongTensor,
+    mask: torch.BoolTensor,
+    filter_value: float = -99999.9999,
+):
+    """Process bad words."""
+    filtered_scores = scores.gather(1, bad_words)
+    filtered_scores = mask.to(filtered_scores.dtype) * filter_value + filtered_scores
+    scores.scatter_(1, bad_words, filtered_scores)
+    return scores
+
+
+logits_process._process_bad_words_ = _process_bad_words_
+
+########## below is for ascend310P ##########
+
+if SocVersion.is_Ascend310P():
+    # Layz import for Ascend310P
+    import torch.distributed as dist
+    from lmdeploy.utils import get_logger
+    from lmdeploy.pytorch.distributed import get_dist_manager, DistContext
+    from lmdeploy.pytorch.engine.model_agent import (
+        msg_with_rank,
+        BaseModelAgent,
+    )
+    from lmdeploy.pytorch.engine.cache_engine import CacheEngine
+    from lmdeploy.pytorch.models.patch import (
+        update_custom_module_map,
+        build_patched_model,
+        add_adapters,
+    )
+    from lmdeploy.pytorch.weight_loader.model_weight_loader import ModelWeightLoader
+    from lmdeploy.pytorch.disagg.config import EngineRole
+
+    logger = get_logger("lmdeploy")
+
+    def _broadcast_next_token_310P(
+        self, next_token_ids: torch.Tensor, dist_ctx: DistContext = None
+    ):
+        # NOTE: Ascend310P does not support broadcast, so we use need to use gloo for broadcast next_token_ids and then transfer it to npu
+        # This mock for properly broadcasting next_token_ids on Ascend 310P device.
+        if dist_ctx is None:
+            dist_ctx = get_dist_manager().current_context()
+        if self.cache_config.role == EngineRole.Decode:
+            next_token_ids = next_token_ids.cpu()
+            tp_cpu_group = dist_ctx.tp_cpu_group
+            dist.all_reduce(next_token_ids, op=dist.ReduceOp.SUM, group=tp_cpu_group)
+        else:
+            # NOTE: Ascend310P does not support broadcast, so we use need to use gloo for broadcast next_token_ids and then transfer it to npu
+            tp_cpu_group = dist_ctx.tp_cpu_group
+            original_device = next_token_ids.device
+            next_token_ids = next_token_ids.cpu()
+            dist.broadcast(next_token_ids, src=0, group=tp_cpu_group)
+            next_token_ids = next_token_ids.to(original_device)
+        return next_token_ids
+
+    def _allocate_cache_310P(self, num_blocks: int, device: torch.device):
+        """
+        allocate cache implement.
+        # NOTE. Ascend300I duo devices require kv_cache to be acl NZ format.
+        """
+        key_block_shape = self.get_key_block_shape(local=True)
+        value_block_shape = self.get_value_block_shape(local=True)
+
+        num_layers = self.num_layers
+        kv_cache_dtype = self.kv_cache_dtype
+
+        if device != "cpu":
+            import torch_npu
+
+            key_cache = torch_npu.empty_with_format(
+                size=(num_layers, num_blocks, *key_block_shape),
+                dtype=kv_cache_dtype,
+                device="npu",
+                acl_format=29,  # 29 for acl NZ format
+            )
+            value_cache = torch_npu.empty_with_format(
+                size=(num_layers, num_blocks, *value_block_shape),
+                dtype=kv_cache_dtype,
+                device="npu",
+                acl_format=29,
+            )
+        else:
+            key_cache = torch.empty(
+                size=(num_layers, num_blocks, *key_block_shape),
+                dtype=kv_cache_dtype,
+                device=device,
+            )
+            value_cache = torch.empty(
+                size=(num_layers, num_blocks, *value_block_shape),
+                dtype=kv_cache_dtype,
+                device=device,
+            )
+
+        output = (key_cache, value_cache)
+
+        if self.cache_config.quant_policy in (4, 8):
+            dtype = self.model_config.dtype
+            key_sz_cache = torch.empty(
+                size=(num_layers, num_blocks, *key_block_shape[:-1], 2),
+                dtype=dtype,
+                device=device,
+            )
+            val_sz_cache = torch.empty(
+                size=(num_layers, num_blocks, *value_block_shape[:-1], 2),
+                dtype=dtype,
+                device=device,
+            )
+            output = output + (key_sz_cache, val_sz_cache)
+
+        return output
+
+    @torch.inference_mode()
+    def load_model_weights_310P(
+        model: torch.nn.Module,
+        checkpoint_path: str,
+        prefix: str = None,
+        device: torch.device = None,
+    ):
+        """Loading model weights."""
+        loader = ModelWeightLoader(checkpoint_path, prefix=prefix)
+        loader.load_model_weights(model, device=device)
+        model.eval()
+        # NOTE: Ascend310P convert Linear weight to NZ format defaultly in graph mode.
+        # However, vision_model part is not compiled in graph mode, so we skip converting weights of vision_model part.
+        # This is a workaround for Ascend310P.
+        for name, mod in model.named_modules():
+            if (
+                not hasattr(mod, "update_weights")
+                or name.startswith("vision_model")
+                or name.startswith("visual")
+            ):
+                continue
+            mod.update_weights()
+
+    def _build_model_310P(self):
+        """
+        Build patched model.
+        NOTE: Ascend310P convert Linear weight to NZ format defaultly in graph mode.
+        However, vision_model part is not compiled in graph mode, so we skip converting weights of vision_model part.
+        """
+        model_path = self.model_path
+        adapters = self.adapters
+        device = self.device
+        rank = self.rank
+        custom_module_map = self.model_config.custom_module_map
+        if custom_module_map is not None:
+            update_custom_module_map(custom_module_map)
+        logger.debug(msg_with_rank(rank, "build model."))
+        patched_model = build_patched_model(
+            self.model_config, device=device, model_format=self.misc_config.model_format
+        )
+        logger.debug(msg_with_rank(rank, "loading weights."))
+        if not self.misc_config.empty_init:
+            load_model_weights_310P(patched_model, model_path, device=device)
+        if adapters is not None:
+            logger.debug(msg_with_rank(rank, "loading adapters."))
+            add_adapters(
+                patched_model, adapters, dtype=self.model_config.dtype, device=device
+            )
+        self.patched_model = patched_model
+
+    # Ascend310P dose't support broadcast for now, so we need to use gloo for broadcast next_token_ids and then transfer it to npu
+    BaseModelAgent._broadcast_next_token = _broadcast_next_token_310P
+    # Ascend310P requires kv_cache to be acl NZ format. So allocate gpu cache in NZ format.
+    CacheEngine._allocate_cache = _allocate_cache_310P
+    # We convert Linear weight to NZ format on Ascend310P device defaultly in graph mode.
+    # However, vision_model part is not compiled in graph mode, so we skip converting weights of vision_model part.
+    BaseModelAgent._build_model = _build_model_310P

dlinfer/framework/lmdeploy_ext/device/camb.py

@@ -0,0 +1,24 @@
+# Copyright (c) 2024, DeepLink. All rights reserved.
+import torch
+
+from lmdeploy.pytorch.backends.default.multinomial_sampling import (
+    DefaultMultinomialSamplingImpl,
+)
+
+
+def CambDefaultMultinomialSamplingImpl_forward(
+    self,
+    scores: torch.Tensor,
+    seeds: torch.LongTensor,
+    offsets: torch.LongTensor,
+    indices: torch.Tensor = None,
+):
+    r"""
+    Note.torch_mlu.multinomial dosen't support replacement=True, whereas lmdeploy set replacement=True by default.
+    """
+    sampled_index = torch.multinomial(scores, num_samples=1, replacement=False)
+    outputs = torch.gather(indices, dim=1, index=sampled_index)
+    return outputs.view(-1)
+
+
+DefaultMultinomialSamplingImpl.forward = CambDefaultMultinomialSamplingImpl_forward

dlinfer/framework/lmdeploy_ext/quants/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) 2024, DeepLink. All rights reserved.
+import importlib
+from functools import lru_cache
+from dlinfer.vendor import vendor_name
+
+
+awq_vendor = ["ascend"]
+
+
+@lru_cache(1)
+def import_vendor_module(vendor_name_str):
+    if vendor_name_str in awq_vendor:
+        importlib.import_module(f".{vendor_name_str}_awq", __package__)
+
+
+def vendor_quant_init():
+    import_vendor_module(vendor_name)
+
+
+vendor_quant_init()