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

dlinfer/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) 2024, DeepLink. All rights reserved.
+import dlinfer.vendor as vendor
+
+vendor.vendor_torch_init()
+__version__ = "0.2.3.post2"

dlinfer/framework/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) 2024, DeepLink. All rights reserved.

dlinfer/framework/lmdeploy_ext/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) 2024, DeepLink. All rights reserved.
+import dlinfer.framework.transformers_ext
+import dlinfer.framework.torch_npu_ext
+from . import quants
+from . import cudagraph
+from . import device

dlinfer/framework/lmdeploy_ext/cudagraph/__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
+
+
+graph_vendor = ["ascend", "maca", "camb", "ppu"]
+
+
+@lru_cache(1)
+def import_vendor_module(vendor_name_str):
+    if vendor_name_str in graph_vendor:
+        importlib.import_module(f".{vendor_name_str}_cudagraph", __package__)
+
+
+def vendor_graph_init():
+    import_vendor_module(vendor_name)
+
+
+vendor_graph_init()

dlinfer/framework/lmdeploy_ext/cudagraph/ascend_cudagraph.py

@@ -0,0 +1,391 @@
+# Copyright (c) 2024, OpenMMLab and DeepLink. All rights reserved.
+# this file implements the cudagraph for ascend backend.
+import functools
+from typing import Any, Dict, List
+from contextlib import ExitStack
+import torch
+from torch import Tensor
+from torch.profiler import record_function
+
+from lmdeploy.pytorch.models.utils.cudagraph import CudaGraphMeta
+from lmdeploy.pytorch.models.utils.cudagraph import CudaGraphMixin
+from lmdeploy.pytorch.config import BackendConfig, CacheConfig, ModelConfig
+from lmdeploy.pytorch.model_inputs import StepContext, get_step_ctx_manager
+from lmdeploy.pytorch.backends.graph_runner import GraphRunner
+from lmdeploy.pytorch.backends.cuda import graph_runner
+
+from lmdeploy.utils import get_logger
+
+logger = get_logger("dlinfer")
+BuffType = Dict[str, Tensor]
+
+
+# AscendCudaGraphMixin methods for cudagraph buffer management.
+def AscendCudaGraphMixin_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.empty(
+        1, max_tokens, dtype=torch.int32, device=device
+    )
+
+    input_buffers["position_ids"] = torch.empty(
+        (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.ones(
+        max_batches, dtype=torch.int32, device=device
+    )
+
+    input_buffers["kv_seqlens"] = [0] * max_batches
+
+    input_buffers["q_start_loc"] = torch.arange(
+        max_batches + 1, dtype=torch.int32, device=device
+    )
+
+    input_buffers["kv_start_indices"] = -torch.ones(
+        (max_batches), dtype=torch.int64, device=device
+    )
+    return input_buffers
+
+
+def AscendCudaGraphMixin_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
+    kv_seqlens: List = 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["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 = get_ascend_compatible_size(batch_size)
+
+    attn_metadata.block_offsets = input_buffers["block_offsets"][: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]
+
+    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 AscendCudaGraphMixin_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 = AscendCudaGraphMixin_make_buffers_cudagraph
+CudaGraphMixin.fill_buffers_cudagraph = AscendCudaGraphMixin_fill_buffers_cudagraph
+CudaGraphMixin.update_context_cudagraph = AscendCudaGraphMixin_update_context_cudagraph
+
+
+def next_power_of_2(n: int):
+    """Return the smallest power of 2 greater than or equal to n."""
+    n -= 1
+    n |= n >> 1
+    n |= n >> 2
+    n |= n >> 4
+    n |= n >> 8
+    n |= n >> 16
+    n |= n >> 32
+    n += 1
+    return n
+
+
+def get_ascend_compatible_size(n: int):
+    """Get ascend compatible size."""
+    if n <= 16:
+        n = next_power_of_2(n)
+    elif n <= 256:
+        n = (n + 15) & ~0xF
+    else:
+        n = (((n - 1) >> 8) + 1) << 8
+    return n
+
+
+@functools.lru_cache
+def _get_capture_batch_size_impl(max_batches: int):
+    """Capture batch size."""
+    ret = []
+    batch_size = 1
+    batch_step_1, batch_step_2 = 16, 256
+    # power of 2
+    while batch_size <= min(batch_step_1, max_batches):
+        ret.append(batch_size)
+        batch_size *= 2
+
+    # step 1
+    ret += list(range(batch_size, min(max_batches, batch_step_2) + 1, batch_step_1))
+
+    # step 2
+    ret += list(range(ret[-1] + batch_step_2, max_batches + 1, batch_step_2))
+
+    # ensure max_batches in ret
+    if max_batches != ret[-1]:
+        ret.append(max_batches)
+
+    return ret
+
+
+def _false(*args, **kwargs):
+    """Default value of not support cuda graph."""
+    return False
+
+
+class AscendSingleGraphRunner:
+    """Cuda single graph runner."""
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        max_batches: int,
+        max_tokens: int,
+        num_blocks: int,
+        is_decoding: bool,
+        pool: Any,
+        model_config: ModelConfig,
+        device: torch.device,
+    ):
+        self.model = model
+        self.ctx_mgr = model.ctx_mgr
+        self.model_config = model_config
+
+        self.meta = CudaGraphMeta(
+            max_batchs=max_batches,
+            max_tokens=max_tokens,
+            num_blocks=num_blocks,
+            is_decoding=is_decoding,
+            device=device,
+            input_buffers=dict(),
+            output_buffers=dict(),
+            vocab_size=self.model_config.vocab_size,
+        )
+        self.device = device
+        self.max_batches = max_batches
+        self.max_tokens = max_tokens
+        self.num_blocks = num_blocks
+        self.is_decoding = is_decoding
+        self.pool = pool
+        self._graph: torch.npu.NPUGraph = None
+
+    @record_function("capture_cudagraph")
+    def capture(self, **kwargs):
+        """Capture graph."""
+        logger.debug(f"Capturing graph with meta: {self.meta}")
+        self.meta.input_buffers = self.model.make_buffers_cudagraph(self.meta, **kwargs)
+        padded_kwargs = self.model.fill_buffers_cudagraph(self.meta, **kwargs)
+        context = self.ctx_mgr.current_context()
+        self.model.update_context_cudagraph(self.meta, context)
+        current_stream = torch.cuda.current_stream()
+
+        aclgraph = torch.npu.NPUGraph()
+        with ExitStack() as stack:
+            with torch.npu.graph(
+                aclgraph,
+                auto_dispatch_capture=True,
+                pool=self.pool,
+                stream=current_stream,
+            ):
+                output = self.model(**padded_kwargs)
+
+        output_buffers = dict(logits=output)
+        self.meta.output_buffers = output_buffers
+        self._graph = aclgraph
+        return output
+
+    @record_function("forward_cudagraph")
+    def forward(self, **kwargs):
+        """forward."""
+        num_tokens = kwargs["input_ids"].size(-1)
+        assert self._graph is not None
+        self.model.fill_buffers_cudagraph(self.meta, **kwargs)
+        context = self.ctx_mgr.current_context()
+        self.model.update_context_cudagraph(self.meta, context)
+        torch.npu.synchronize()
+        self._graph.replay()
+        self._graph.update(
+            cpu_update_input=[
+                {"actual_seq_lengths_kv": self.meta.input_buffers["kv_seqlens"]}
+            ]
+        )
+
+        output = self.meta.output_buffers["logits"][:, :num_tokens]
+        return output
+
+    def __del__(self):
+        """del."""
+        del self._graph
+
+
+class AscendGraphRunner(GraphRunner):
+    """Cuda graph runner."""
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        model_config: ModelConfig,
+        cache_config: CacheConfig,
+        backend_config: BackendConfig,
+        device: torch.device,
+    ):
+        super().__init__(model, model_config, cache_config, backend_config, device)
+        self.max_batches = cache_config.max_batches
+        self.max_tokens = cache_config.max_prefill_token_num
+        self.num_blocks = cache_config.num_gpu_blocks
+        self.enable_graph = self.check_enable_graph()
+        self.graph_pool_handle = torch.cuda.graph_pool_handle()
+        self._runner_map: Dict[Any, AscendSingleGraphRunner] = dict()
+        self.has_try_compile_model: bool = False
+
+    def check_enable_graph(self):
+        """Check enable graph."""
+        if self.backend_config.eager_mode:
+            return _false
+
+        return getattr(self.model, "support_cuda_graph", _false)
+
+    def _get_capture_tokens(self, batch_size: int):
+        """Get capture tokens."""
+        cap_sizes = self.get_capture_batch_sizes()
+        for size in cap_sizes:
+            if size >= batch_size:
+                return size
+        assert False, f"Unsupported batch_size={batch_size}"
+
+    def get_graph_key(
+        self,
+        input_ids: torch.Tensor,
+        **kwargs,
+    ):
+        """Get graph key."""
+        context = self.ctx_mgr.current_context()
+        is_decoding = context.is_decoding
+        num_tokens = input_ids.numel()
+        meta = self.get_meta()
+        enable_microbatch = get_step_ctx_manager().current_context().enable_microbatch
+        if meta.padding_batch_size is None:
+            new_num_tokens = self._get_capture_tokens(num_tokens)
+        else:
+            new_num_tokens = self._get_capture_tokens(meta.padding_batch_size)
+        return (new_num_tokens, is_decoding, enable_microbatch)
+
+    def __call__(self, **kwargs):
+        """call."""
+        enable_graph = self.enable_graph(**kwargs)
+
+        if not enable_graph:
+            with record_function("forward_eager"):
+                ret = self.model(**kwargs)
+                return ret
+
+        graph_key = self.get_graph_key(**kwargs)
+        max_tokens = graph_key[0]
+        is_decoding = graph_key[1]
+        if graph_key not in self._runner_map:
+            max_batches = max_tokens if is_decoding else self.max_batches
+            runner = AscendSingleGraphRunner(
+                self.model,
+                max_batches=max_batches,
+                max_tokens=max_tokens,
+                num_blocks=self.num_blocks,
+                is_decoding=is_decoding,
+                pool=self.graph_pool_handle,
+                model_config=self.model_config,
+                device=self.device,
+            )
+            runner.capture(**kwargs)
+            self._runner_map[graph_key] = runner
+        else:
+            runner = self._runner_map[graph_key]
+        output = runner.forward(**kwargs)
+        return output
+
+    @record_function("prepare_inputs_for_generation")
+    def prepare_inputs_for_generation(
+        self,
+        past_key_values: List[List[torch.Tensor]],
+        inputs_embeds: torch.Tensor = None,
+        context: StepContext = None,
+    ):
+        """Prepare inputs."""
+        return self.model.prepare_inputs_for_generation(
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            context=context,
+        )
+
+    def reset(self):
+        """Remove all graphs to prevent hanging on exit."""
+        self._runner_map.clear()
+
+    def update_inputs(self, inputs):
+        """Update inputs."""
+        if self.backend_config.eager_mode:
+            return inputs
+        is_decoding = inputs.is_decoding
+        dp_meta = inputs.dp_meta
+        if is_decoding and dp_meta is not None:
+            meta = self.get_meta()
+            padding_batch_size = meta.padding_batch_size
+            tp_size = self._get_capture_tokens(padding_batch_size)
+            dp_meta.tp_sizes = [tp_size] * len(dp_meta.tp_sizes)
+        return inputs
+
+    def get_capture_batch_sizes(self) -> List[int]:
+        """Capture batch sizes."""
+        return _get_capture_batch_size_impl(self.cache_config.max_batches)
+
+
+graph_runner.CUDAGraphRunner = AscendGraphRunner

dlinfer/framework/lmdeploy_ext/cudagraph/camb_cudagraph.py

@@ -0,0 +1,133 @@
+# Copyright (c) 2024, 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 CambCudaGraphMixin_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.ones(
+        (1, max_tokens), dtype=torch.int32, device=device
+    )
+
+    input_buffers["block_offsets"] = torch.zeros(
+        (max_batches, num_blocks), dtype=torch.int32, device=device
+    )
+
+    # in attn_metadata, we mock q_start_loc as cu_seqlens so need to add 1
+    input_buffers["q_start_loc"] = torch.arange(
+        max_batches + 1, dtype=torch.int32, device=device
+    )
+
+    input_buffers["q_seqlens"] = torch.ones(
+        max_batches, dtype=torch.int32, device=device
+    )
+
+    input_buffers["kv_seqlens"] = torch.ones(
+        max_batches, dtype=torch.int32, device=device
+    )
+    # critical to set negative for kv_start_indices
+    # if we don't set it, two batches with same input tokens
+    # will result in different answer
+    input_buffers["kv_start_indices"] = -torch.ones(
+        max_batches, dtype=torch.int32, device=device
+    )
+
+    return input_buffers
+
+
+def CambCudaGraphMixin_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_seqlens"][:batch_size] = q_seqlens
+    input_buffers["kv_seqlens"][:batch_size] = kv_seqlens
+    input_buffers["q_start_loc"][: batch_size + 1] = q_start_loc
+    input_buffers["kv_start_indices"][:num_tokens] = kv_start_indices[:num_tokens]
+
+    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
+
+    # below only used for capture graph
+    # create inputs
+    new_num_tokens = next_power_of_2(num_tokens)
+    new_batch_size = new_num_tokens
+
+    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_num_tokens]
+    new_inputs = dict(
+        past_key_values=past_key_values,
+        attn_metadata=attn_metadata,
+    )
+
+    # is_decoding:
+    new_inputs["input_ids"] = input_buffers["input_ids"][:, : new_batch_size + 1]
+    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 CambCudaGraphMixin_update_context_cudagraph(self, graph_meta, context):
+    """update step context with input buffers."""
+    input_buffers = graph_meta.input_buffers
+    context.q_seqlens = input_buffers["q_seqlens"]
+    context.kv_seqlens = input_buffers["kv_seqlens"]
+    context.q_start_loc = input_buffers["q_start_loc"]
+
+
+CudaGraphMixin.make_buffers_cudagraph = CambCudaGraphMixin_make_buffers_cudagraph
+CudaGraphMixin.fill_buffers_cudagraph = CambCudaGraphMixin_fill_buffers_cudagraph
+CudaGraphMixin.update_context_cudagraph = CambCudaGraphMixin_update_context_cudagraph

dlinfer/framework/lmdeploy_ext/cudagraph/maca_cudagraph.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2024, 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 MacaCudaGraphMixin_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.empty(
+        1, max_tokens, dtype=torch.int32, device=device
+    )
+
+    input_buffers["position_ids"] = torch.empty(
+        (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.ones(
+        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.arange(
+        max_batches + 1, dtype=torch.int32, device=device
+    )
+
+    input_buffers["kv_start_indices"] = -torch.ones(
+        (max_batches, 1), dtype=torch.int64, device=device
+    )
+    return input_buffers
+
+
+def MacaCudaGraphMixin_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]
+
+    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 MacaCudaGraphMixin_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 = MacaCudaGraphMixin_make_buffers_cudagraph
+CudaGraphMixin.fill_buffers_cudagraph = MacaCudaGraphMixin_fill_buffers_cudagraph
+CudaGraphMixin.update_context_cudagraph = MacaCudaGraphMixin_update_context_cudagraph