def test_gpt2_multiple_token_generation(model_name, optimized):
    """Generation when we pass in multiple tokens at a time, not just one."""
    dtype = torch.float16
    device = "cuda"
    rtol, atol = 3e-3, 3e-1
    config = GPT2Config.from_pretrained(model_name)
    config.residual_in_fp32 = True
    if optimized:
        config.use_flash_attn = True
        config.fused_bias_fc = True
        config.fused_mlp = True
        config.fused_dropout_add_ln = True

    model = GPTLMHeadModel.from_pretrained(model_name, config, device=device, dtype=dtype)
    model.eval()

    torch.manual_seed(0)
    input_ids = torch.randint(0, config.vocab_size, (1, 20), dtype=torch.long, device=device)
    # Reference logits
    logits_ref = model(input_ids).logits

    # Run 10 tokens, then pass in another 4, then another 6, to see if we get the same logits
    inference_params = InferenceParams(max_seqlen=20, max_batch_size=1)
    logits_10 = model(input_ids[:, :10], inference_params=inference_params).logits
    inference_params.seqlen_offset += 10
    position_ids = torch.arange(10, 14, dtype=torch.long, device=device)
    logits_1014 = model(
        input_ids[:, 10:14], position_ids=position_ids, inference_params=inference_params
    ).logits
    inference_params.seqlen_offset += 4
    position_ids = torch.arange(14, 20, dtype=torch.long, device=device)
    logits_1420 = model(
        input_ids[:, 14:20], position_ids=position_ids, inference_params=inference_params
    ).logits
    logits = torch.cat([logits_10, logits_1014, logits_1420], dim=1)
    print(f"Logits max diff: {(logits - logits_ref).abs().max().item()}")
    print(f"Logits mean diff: {(logits - logits_ref).abs().mean().item()}")
    assert torch.allclose(logits, logits_ref, rtol=rtol, atol=atol)
Generation when we pass in multiple tokens at a time, not just one.
tests/models/test_gpt.py