Function check_streamingllm

tests/test_infer/test_streamingllm.py:27–100

()

def check_streamingllm():
    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_key_value_heads=2,
            num_hidden_layers=16,
        )
    ).cuda()
    model = model.eval()

    input_token_ids = data_gen(1, 4)

    output_len = 128

    inference_config = InferenceConfig(
        max_batch_size=1,
        max_output_len=output_len,
        dtype="fp32",
        use_cuda_kernel=True,
        enable_streamingllm=True,
        start_token_size=4,
        generated_token_size=32,
    )

    inference_engine = InferenceEngine(model, tokenizer, inference_config, verbose=True)
    assert inference_engine.generation_config.max_new_tokens == output_len
    inference_engine.add_request(prompts_token_ids=input_token_ids)
    assert inference_engine.request_handler._has_waiting()

    assert inference_config.start_token_size == inference_config.block_size

    request_handler = inference_engine.request_handler
    running_bb = request_handler.running_bb

    for _ in range(12):
        inference_engine.step()

    assert running_bb.block_tables[0].tolist() == [0, -1, -1, -1]
    assert running_bb.seq_lengths[0].item() == 16

    for _ in range(16):
        inference_engine.step()

    assert running_bb.block_tables[0].tolist() == [0, 1, -1, -1]
    assert running_bb.seq_lengths[0].item() == 32

    for _ in range(16):
        inference_engine.step()

    assert running_bb.block_tables[0].tolist() == [0, 1, 2, -1]
    assert running_bb.seq_lengths[0].item() == 48

    for _ in range(16):