code: def test_gpt2_generation_cg(model_name, rotary, seqlen, maxlen):
      """Check that decoding with CUDA graph is the same as decoding without CUDA graph."""
      dtype = torch.float16
      device = "cuda"
      rtol, atol = 3e-3, 3e-1
      config = GPT2Config.from_pretrained(model_name)
      config.n_positions = 16 * 1024
      assert seqlen <= maxlen <= config.n_positions
      if rotary is not None:
          config.n_positions = 0
          config.rotary_emb_dim = 32
          config.rotary_emb_interleaved = rotary == "interleaved"
      config.residual_in_fp32 = True
      config.use_flash_attn = True
      config.fused_bias_fc = True
      config.fused_mlp = True
      config.fused_dropout_add_ln = True
  
      model = GPTLMHeadModel(config, device=device, dtype=dtype)
      model.eval()
  
      torch.manual_seed(0)
      batch_size = 1
      input_ids = torch.randint(
          0,
          config.vocab_size,
          (batch_size, seqlen),
          dtype=torch.long,
          device=device,
      )
      teacher_outputs = torch.randint(
          0,
          config.vocab_size,
          (batch_size, maxlen),
          dtype=torch.long,
          device=device,
      )
  
      logits = get_logits(
          model, input_ids, maxlen, teacher_outputs=teacher_outputs
      )
      logits_cg = get_logits(
          model,
          input_ids,
          maxlen,
          teacher_outputs=teacher_outputs,
          cg=True,
      )
      assert torch.equal(logits, logits_cg)
  
      # Try increasing batch size and seqlen, then decrease them to see if it's still correct
      batch_size = 3
      maxlen += 30
      input_ids = torch.randint(
          0,
          config.vocab_size,
          (batch_size, seqlen),
          dtype=torch.long,
          device=device,
      )
      teacher_outputs = torch.randint(
          0,
          config.vocab_size,
          (batch_size, maxlen),
          dtype=torch.long,
          device=device,
      )
      logits = get_logits(
          model, input_ids, maxlen, teacher_outputs=teacher_outputs
      )
      logits_cg = get_logits(
          model,
          input_ids,
          maxlen,
          teacher_outputs=teacher_outputs,
          cg=True,
      )
      assert torch.equal(logits, logits_cg)
  
      batch_size = 2
      maxlen -= 35
      input_ids = torch.randint(
          0,
          config.vocab_size,
          (batch_size, seqlen),
          dtype=torch.long,
          device=device,
      )
      teacher_outputs = torch.randint(
          0,
          config.vocab_size,
          (batch_size, maxlen),
          dtype=torch.long,
          device=device,
      )
      logits = get_logits(
          model, input_ids, maxlen, teacher_outputs=teacher_outputs
      )
      logits_cg = get_logits(
          model,
          input_ids,
          maxlen,
          teacher_outputs=teacher_outputs,
          cg=True,
      )
      assert torch.equal(logits, logits_cg)
  
file_path: flash-attention/tests/models/test_gpt.py
imports: import re
  import pytest
  import torch
  from einops import rearrange
  from flash_attn.models.gpt import GPTLMHeadModel, remap_state_dict_hf_gpt2, shard_state_dict_tp, combine_state_dicts_tp
  from flash_attn.utils.generation import InferenceParams
  from flash_attn.utils.pretrained import state_dict_from_pretrained
  from transformers import GPT2Config, GPT2Tokenizer
  from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
label_desc: Use image node_id 6 for calling a global function with example usage: test_gpt2_generation_cg(model_name, rotary, seqlen, maxlen) without return types