Function check_inference_engine

tests/test_infer/test_cuda_graph.py:21–71

(use_cuda_graph=False, batch_size=32)

def check_inference_engine(use_cuda_graph=False, batch_size=32):
    setup_seed(20)
    tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
    model = (
        LlamaForCausalLM(
            LlamaConfig(
                vocab_size=50000, hidden_size=512, intermediate_size=1536, num_attention_heads=4, num_hidden_layers=16
            )
        )
        .cuda()
        .half()
    )
    model = model.eval()

    prompts_token_ids = []
    for i in range(batch_size):
        prompts_token_ids.append(
            np.random.randint(low=0, high=100, size=random.randint(1, max(1024 // batch_size, 32))).tolist()
        )

    input_len = 1024
    output_len = 128
    do_sample = False
    top_p = 0.5
    top_k = 50

    if use_cuda_graph:
        inference_config = InferenceConfig(
            max_batch_size=batch_size,
            max_input_len=input_len,
            max_output_len=output_len,
            use_cuda_kernel=False,
            use_cuda_graph=True,
            block_size=16,
        )
    else:
        inference_config = InferenceConfig(
            max_batch_size=batch_size,
            max_input_len=input_len,
            max_output_len=output_len,
            use_cuda_kernel=False,
            use_cuda_graph=False,
            block_size=16,
        )

    inference_engine = InferenceEngine(model, tokenizer, inference_config, verbose=True)
    assert inference_engine.generation_config.max_new_tokens == output_len
    generation_config = GenerationConfig(do_sample=do_sample, top_p=top_p, top_k=top_k)
    outputs = inference_engine.generate(prompts_token_ids=prompts_token_ids, generation_config=generation_config)

    return outputs


def check_output_consistency(batch_size):