code: def capture_graph(
      model,
      inference_params,
      batch_size,
      max_seqlen,
      decoding_seqlen=1,
      mempool=None,
      n_warmups=2,
  ):
      device = next(iter(model.parameters())).device
      input_ids = torch.full(
          (batch_size, decoding_seqlen),
          0,
          dtype=torch.long,
          device=device,
      )
      position_ids = torch.full(
          (batch_size, decoding_seqlen),
          0,
          dtype=torch.long,
          device=device,
      )
      seqlen_offset_og = inference_params.seqlen_offset
      inference_params.seqlen_offset = max_seqlen - decoding_seqlen
      inference_params.lengths_per_sample[
          :
      ] = inference_params.seqlen_offset
  
      # Warmup before capture
      s = torch.cuda.Stream()
      s.wait_stream(torch.cuda.current_stream())
      with torch.cuda.stream(s):
          for _ in range(n_warmups):
              logits = model(
                  input_ids,
                  position_ids=position_ids,
                  inference_params=inference_params,
                  num_last_tokens=decoding_seqlen,
              ).logits
          s.synchronize()
          # This might be needed for correctness if we run with NCCL_GRAPH_MIXING_SUPPORT=0,
          # which requires that graph launch and non-captured launch to not overlap (I think,
          # that's how I interpret the documentation). I'm not sure if this is required.
          if torch.distributed.is_initialized():
              torch.distributed.barrier()
      torch.cuda.current_stream().wait_stream(s)
      # Captures the graph
      # To allow capture, automatically sets a side stream as the current stream in the context
      graph = torch.cuda.CUDAGraph()
      with torch.cuda.graph(graph, pool=mempool):
          logits = model(
              input_ids,
              position_ids=position_ids,
              inference_params=inference_params,
              num_last_tokens=decoding_seqlen,
          ).logits
  
      def run(new_input_ids, new_position_ids, seqlen):
          inference_params.lengths_per_sample[:] = seqlen
          input_ids.copy_(new_input_ids)
          position_ids.copy_(new_position_ids)
          graph.replay()
          return logits.clone()
  
      inference_params.seqlen_offset = seqlen_offset_og
      return run
  
imports: import gc
  import time
  from collections import namedtuple
  from dataclasses import dataclass, field
  from functools import partial
  from typing import Callable, Optional, Sequence, Union
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from torch import Tensor
  from torch.profiler import ProfilerActivity, profile, record_function
  from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput
label_desc: Use image node_id 9 for calling a global function with example usage: capture_graph(model, inference_params, batch_size, max_seqlen, decoding_seqlen, mempool, n_warmups) and returns: run, logits