code: def test_llama_generation(model_name, checkpoint_format):
      checkpoint_path = (
          Path(
              os.environ.get(
                  "CHECKPOINT_DIR",
                  current_dir.parent.parent / "checkpoints",
              )
          )
          / "llama"
      )
  
      dtype = torch.float16
      device = "cuda"
      config = config_from_checkpoint(
          checkpoint_path, model_name, checkpoint_format
      )
      config = llama_config_to_gpt2_config(config)
      config.use_flash_attn = True
      config.fused_bias_fc = True
      config.fused_mlp = False  # We don't have fused GatedMLP yet
      config.fused_dropout_add_ln = True
      config.residual_in_fp32 = True
  
      tokenizer = LlamaTokenizer.from_pretrained(
          Path(checkpoint_path) / f"{model_name}-hf"
      )
      eos_token_id = tokenizer.eos_token_id
  
      torch.manual_seed(0)
      batch_size = 1
      seqlen = 100
      max_length = 150
      input_ids = torch.randint(
          0,
          config.vocab_size,
          (batch_size, seqlen),
          dtype=torch.long,
          device=device,
      )
  
      model_hf = LlamaForCausalLM.from_pretrained(
          Path(checkpoint_path) / f"{model_name}-hf",
          torch_dtype=dtype,
          device_map={"": device},
      )
      model_hf.eval()
      print("HF fp16")
      torch.cuda.synchronize()
      start = time.time()
      out_hf = model_hf.generate(
          input_ids=input_ids,
          max_length=max_length,
          return_dict_in_generate=True,
          output_scores=True,
      )
      torch.cuda.synchronize()
      print(
          f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms"
      )
      del model_hf
  
      # Need auto here since the 13B fp32 model doesn't fit in memory on a A100 40GB
      model_ref = LlamaForCausalLM.from_pretrained(
          Path(checkpoint_path) / f"{model_name}-hf", device_map="auto"
      )
      model_ref.eval()
      with torch.no_grad():
          logits_ref = (
              model_ref(out_hf.sequences)
              .logits[:, (seqlen - 1) : -1]
              .to(device=device)
          )
      del model_ref
  
      pretrained_state_dict = _pretrained_state_dict_from_checkpoint(
          checkpoint_path, model_name, config, checkpoint_format
      )
      model = GPTLMHeadModel(config, device=device, dtype=dtype)
      model.load_state_dict(pretrained_state_dict)
      model.eval()
  
      print("Without CUDA graph")
      torch.cuda.synchronize()
      start = time.time()
      out = model.generate(
          input_ids=input_ids,
          max_length=max_length,
          eos_token_id=eos_token_id,
          return_dict_in_generate=True,
          output_scores=True,
          enable_timing=True,
          teacher_outputs=out_hf.sequences,
      )
      torch.cuda.synchronize()
      print(
          f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms"
      )
  
      # Capture graph outside the timing loop
      batch_size, seqlen_og = input_ids.shape
      model._decoding_cache = update_graph_cache(
          model, None, batch_size, seqlen_og, max_length
      )
      print("With CUDA graph")
      torch.cuda.synchronize()
      start = time.time()
      out_cg = model.generate(
          input_ids=input_ids,
          max_length=max_length,
          cg=True,
          return_dict_in_generate=True,
          output_scores=True,
          enable_timing=True,
          teacher_outputs=out_hf.sequences,
      )
      torch.cuda.synchronize()
      print(
          f"Prompt processing + decoding time: {(time.time() - start) * 1000:.0f}ms"
      )
  
      with torch.no_grad():
          logits_parallel = model(out_hf.sequences).logits[
              :, (seqlen - 1) : -1
          ]
      logits_hf = torch.stack(out_hf.scores, dim=1)
      logits = torch.stack(out.scores, dim=1)
      logits_cg = torch.stack(out_cg.scores, dim=1)
  
      del model
  
      hf_error = (logits_hf - logits_ref).abs().max().item()
  
      print(f"HF fp16 logits max diff: {hf_error}")
      print(
          f"Logits max diff: {(logits - logits_ref).abs().max().item()}"
      )
      print(
          f"Logits CG max diff: {(logits_cg - logits_ref).abs().max().item()}"
      )
  
      assert (
          logits_parallel - logits_ref
      ).abs().max().item() < 2 * hf_error
      assert (logits - logits_ref).abs().max().item() < 2 * hf_error
      assert torch.equal(logits_cg, logits)
  
imports: import os
  import time
  from pathlib import Path
  import shutil
  import pytest
  import torch
  from einops import rearrange
  from flash_attn.models.gpt import GPTLMHeadModel, combine_state_dicts_tp, shard_state_dict_tp
  from flash_attn.models.llama import config_from_checkpoint, inv_remap_state_dict_hf_llama, llama_config_to_gpt2_config, remap_state_dict_hf_llama, remap_state_dict_meta_llama, state_dicts_from_checkpoint
  from flash_attn.utils.distributed import all_gather_raw
  from flash_attn.utils.generation import update_graph_cache
  from flash_attn.utils.pretrained import state_dict_from_pretrained
  from transformers import LlamaConfig, LlamaTokenizer
  from transformers.models.llama.modeling_llama import LlamaForCausalLM
  from transformers import AutoConfig
label_desc: Use image node_id 6 for calling a global function with example usage: test_llama_generation(model_name, checkpoint_format) without return types