qwen3-coder-next / code /models /tt_transformers /tests /test_chunked_generation.py

Add files using upload-large-folder tool

babe29f verified 22 days ago

6.85 kB

# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0
import os

import pytest
import torch
from loguru import logger

import ttnn
from models.common.utility_functions import comp_allclose, comp_pcc
from models.tt_transformers.tt.common import PagedAttentionConfig, get_block_size, num_blocks_in_seq
from models.tt_transformers.tt.generator import Generator
from models.tt_transformers.tt.model import Transformer
from models.tt_transformers.tt.model_config import DecodersPrecision, ModelArgs


@torch.no_grad()
@pytest.mark.timeout(900)
@pytest.mark.parametrize(
    "mesh_device",
    [
        {"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
            os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
        )
    ],
    indirect=True,
)
@pytest.mark.parametrize(
    "paged_attention",
    (True,),
    ids=("paged_attention",),
)
@pytest.mark.parametrize(
    "page_params",
    [{"page_block_size": 64, "page_max_num_blocks": 2048}],
)
@pytest.mark.parametrize(
    "seq_len, prefill_chunk_size",
    [(4096, 2048)],
)
@pytest.mark.parametrize(
    "optimizations",
    [
        pytest.param(
            lambda model_args: DecodersPrecision.accuracy(model_args.n_layers, model_args.model_name), id="accuracy"
        ),
    ],
)
@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
def test_chunked_prefill_single_user(
    seq_len,
    prefill_chunk_size,
    paged_attention,
    page_params,
    optimizations,
    mesh_device,
    reset_seeds,
    ensure_gc,
    is_ci_env,
    request,
):
    dtype = ttnn.bfloat8_b
    batch_size = 1  # For prefill we only support batch_size = 1

    # This sets the minimum PCC for each iteration based on optimization mode
    test_id = request.node.callspec.id
    if "accuracy" in test_id:
        pcc = 0.91  # TODO Look on improving PCC
    else:  # performance mode
        assert "performance" in test_id
        pcc = 0.869  # TODO Look on improving PCC

    model_args = ModelArgs(
        mesh_device, max_batch_size=batch_size, optimizations=optimizations, max_seq_len=seq_len, cache_hf=True
    )
    model_args.max_prefill_chunk_size = prefill_chunk_size

    logger.info("Loading weights...")
    state_dict_prefix = model_args.get_state_dict_prefix("", None)
    state_dict = model_args.load_state_dict()
    reference_state_dict = {
        k[len(state_dict_prefix) :]: v
        for k, v in state_dict.items()
        if (
            any([f"{state_dict_prefix}layers.{i}." in k for i in range(model_args.n_layers)])
            or any(
                [
                    f"{state_dict_prefix}{name}" in k
                    for name in ["tok_embeddings.weight", "learnable_embedding.weight", "norm.weight", "output.weight"]
                ]
            )
        )
    }
    logger.info("Finished loading weights...")

    reference_model = model_args.reference_transformer()
    reference_model.load_state_dict(reference_state_dict)
    embd = model_args.reference_embedding()
    embd.load_state_dict({"emb.weight": state_dict[f"{state_dict_prefix}tok_embeddings.weight"]})

    # Setup page table
    paged_attention_config = PagedAttentionConfig(
        block_size=page_params["page_block_size"],
        max_num_blocks=page_params["page_max_num_blocks"],
    )
    # Implied shuffling of blocks
    # Physical block 0 is reserved as null block in vLLM, so use blocks 1 to max_num_blocks-1
    # (permute max_num_blocks-1 values, then add 1 to shift range from 0..max-2 to 1..max-1)
    num_usable_blocks = paged_attention_config.max_num_blocks - 1
    permutation = torch.randperm(num_usable_blocks)
    # Page table which maps virtual blocks to physical (offset by 1 to skip block 0)
    reverse_permutation = torch.argsort(permutation) + 1
    static_page_table = reverse_permutation.reshape(
        model_args.max_batch_size, num_usable_blocks // model_args.max_batch_size
    )

    # Load TTNN model
    tt_model = Transformer(
        args=model_args,
        mesh_device=mesh_device,
        dtype=dtype,
        state_dict=state_dict,
        weight_cache_path=model_args.weight_cache_path(dtype),
        paged_attention_config=paged_attention_config,
    )
    generator = Generator([tt_model], [model_args], mesh_device)

    logger.info("Model and caches loaded.")

    # Select the first token from the prompt for initial decoding
    pt_prefill_input = torch.randint(0, 32000, (batch_size, seq_len), dtype=torch.long)
    tt_prefill_input = pt_prefill_input

    pt_prefill_input = embd(pt_prefill_input).view(batch_size, seq_len, -1)

    tt_kv_cache = [l.attention.layer_past for l in tt_model.layers]
    # Slice out relevant part of page table
    block_size = get_block_size(tt_kv_cache)
    num_blocks = num_blocks_in_seq(seq_len, block_size)
    static_page_table = static_page_table[:, :num_blocks]

    start_pos = 0
    logger.info("Running reference model")
    ref_output = reference_model(pt_prefill_input, start_pos, mode="decode")

    # Run TT model for various last_token_idxs and start_pos values
    # to test the chunked prefill and prefix caching functionalities.
    # These are implemented together, primarily in
    # Generator.prefill_forward_single_user_text(), both using chunked SDPA,
    # and thus tested together here.
    logger.info("Running TT model")
    for last_token_idx in [
        prefill_chunk_size - 2,  # one chunk minus one token
        prefill_chunk_size - 1,  # exactly one chunk
        prefill_chunk_size,  # one chunk plus one token
        prefill_chunk_size + 1,  # one chunk plus two tokens
        seq_len - 10,  # less than seq_len (two chunks)
        seq_len - 1,  # exactly seq_len (two chunks)
    ]:
        prefill_input_trimmed = tt_prefill_input[:, : last_token_idx + 1]

        for start_pos in [
            0,
            1 * block_size,
            2 * block_size,
            3 * block_size,
            4 * block_size,
        ]:  # Reuse zero or more blocks of cache
            logger.info(f"Running TT model for last_token_idx: {last_token_idx}, start_pos: {start_pos}")
            tt_output_torch = generator.prefill_forward_text(
                prefill_input_trimmed,
                page_table=static_page_table,
                kv_cache=[tt_kv_cache],
                enable_trace=False,
                start_pos=[start_pos],
            )
            ref_output_slice = ref_output[:, last_token_idx : last_token_idx + 1, :]

            passing, pcc_message = comp_pcc(ref_output_slice, tt_output_torch, pcc)

            logger.info(comp_allclose(ref_output_slice, tt_output_torch))
            logger.info(
                f"passing: {passing}, PCC: {pcc_message} (for last_token_idx: {last_token_idx}, start_pos: {start_pos})"
            )
            assert passing

	# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.

	# SPDX-License-Identifier: Apache-2.0
	import os

	import pytest
	import torch
	from loguru import logger

	import ttnn
	from models.common.utility_functions import comp_allclose, comp_pcc
	from models.tt_transformers.tt.common import PagedAttentionConfig, get_block_size, num_blocks_in_seq
	from models.tt_transformers.tt.generator import Generator
	from models.tt_transformers.tt.model import Transformer
	from models.tt_transformers.tt.model_config import DecodersPrecision, ModelArgs


	@torch.no_grad()
	@pytest.mark.timeout(900)
	@pytest.mark.parametrize(
	"mesh_device",
	[
	{"N150": (1, 1), "N300": (1, 2), "T3K": (1, 8), "TG": (8, 4)}.get(
	os.environ.get("MESH_DEVICE"), len(ttnn.get_device_ids())
	)
	],
	indirect=True,
	)
	@pytest.mark.parametrize(
	"paged_attention",
	(True,),
	ids=("paged_attention",),
	)
	@pytest.mark.parametrize(
	"page_params",
	[{"page_block_size": 64, "page_max_num_blocks": 2048}],
	)
	@pytest.mark.parametrize(
	"seq_len, prefill_chunk_size",
	[(4096, 2048)],
	)
	@pytest.mark.parametrize(
	"optimizations",
	[
	pytest.param(
	lambda model_args: DecodersPrecision.accuracy(model_args.n_layers, model_args.model_name), id="accuracy"
	),
	],
	)
	@pytest.mark.parametrize("device_params", [{"fabric_config": True}], indirect=True)
	def test_chunked_prefill_single_user(
	seq_len,
	prefill_chunk_size,
	paged_attention,
	page_params,
	optimizations,
	mesh_device,
	reset_seeds,
	ensure_gc,
	is_ci_env,
	request,
	):
	dtype = ttnn.bfloat8_b
	batch_size = 1 # For prefill we only support batch_size = 1

	# This sets the minimum PCC for each iteration based on optimization mode
	test_id = request.node.callspec.id
	if "accuracy" in test_id:
	pcc = 0.91 # TODO Look on improving PCC
	else: # performance mode
	assert "performance" in test_id
	pcc = 0.869 # TODO Look on improving PCC

	model_args = ModelArgs(
	mesh_device, max_batch_size=batch_size, optimizations=optimizations, max_seq_len=seq_len, cache_hf=True
	)
	model_args.max_prefill_chunk_size = prefill_chunk_size

	logger.info("Loading weights...")
	state_dict_prefix = model_args.get_state_dict_prefix("", None)
	state_dict = model_args.load_state_dict()
	reference_state_dict = {
	k[len(state_dict_prefix) :]: v
	for k, v in state_dict.items()
	if (
	any([f"{state_dict_prefix}layers.{i}." in k for i in range(model_args.n_layers)])
	or any(
	[
	f"{state_dict_prefix}{name}" in k
	for name in ["tok_embeddings.weight", "learnable_embedding.weight", "norm.weight", "output.weight"]
	]
	)
	)
	}
	logger.info("Finished loading weights...")

	reference_model = model_args.reference_transformer()
	reference_model.load_state_dict(reference_state_dict)
	embd = model_args.reference_embedding()
	embd.load_state_dict({"emb.weight": state_dict[f"{state_dict_prefix}tok_embeddings.weight"]})

	# Setup page table
	paged_attention_config = PagedAttentionConfig(
	block_size=page_params["page_block_size"],
	max_num_blocks=page_params["page_max_num_blocks"],
	)
	# Implied shuffling of blocks
	# Physical block 0 is reserved as null block in vLLM, so use blocks 1 to max_num_blocks-1
	# (permute max_num_blocks-1 values, then add 1 to shift range from 0..max-2 to 1..max-1)
	num_usable_blocks = paged_attention_config.max_num_blocks - 1
	permutation = torch.randperm(num_usable_blocks)
	# Page table which maps virtual blocks to physical (offset by 1 to skip block 0)
	reverse_permutation = torch.argsort(permutation) + 1
	static_page_table = reverse_permutation.reshape(
	model_args.max_batch_size, num_usable_blocks // model_args.max_batch_size
	)

	# Load TTNN model
	tt_model = Transformer(
	args=model_args,
	mesh_device=mesh_device,
	dtype=dtype,
	state_dict=state_dict,
	weight_cache_path=model_args.weight_cache_path(dtype),
	paged_attention_config=paged_attention_config,
	)
	generator = Generator([tt_model], [model_args], mesh_device)

	logger.info("Model and caches loaded.")

	# Select the first token from the prompt for initial decoding
	pt_prefill_input = torch.randint(0, 32000, (batch_size, seq_len), dtype=torch.long)
	tt_prefill_input = pt_prefill_input

	pt_prefill_input = embd(pt_prefill_input).view(batch_size, seq_len, -1)

	tt_kv_cache = [l.attention.layer_past for l in tt_model.layers]
	# Slice out relevant part of page table
	block_size = get_block_size(tt_kv_cache)
	num_blocks = num_blocks_in_seq(seq_len, block_size)
	static_page_table = static_page_table[:, :num_blocks]

	start_pos = 0
	logger.info("Running reference model")
	ref_output = reference_model(pt_prefill_input, start_pos, mode="decode")

	# Run TT model for various last_token_idxs and start_pos values
	# to test the chunked prefill and prefix caching functionalities.
	# These are implemented together, primarily in
	# Generator.prefill_forward_single_user_text(), both using chunked SDPA,
	# and thus tested together here.
	logger.info("Running TT model")
	for last_token_idx in [
	prefill_chunk_size - 2, # one chunk minus one token
	prefill_chunk_size - 1, # exactly one chunk
	prefill_chunk_size, # one chunk plus one token
	prefill_chunk_size + 1, # one chunk plus two tokens
	seq_len - 10, # less than seq_len (two chunks)
	seq_len - 1, # exactly seq_len (two chunks)
	]:
	prefill_input_trimmed = tt_prefill_input[:, : last_token_idx + 1]

	for start_pos in [
	0,
	1 * block_size,
	2 * block_size,
	3 * block_size,
	4 * block_size,
	]: # Reuse zero or more blocks of cache
	logger.info(f"Running TT model for last_token_idx: {last_token_idx}, start_pos: {start_pos}")
	tt_output_torch = generator.prefill_forward_text(
	prefill_input_trimmed,
	page_table=static_page_table,
	kv_cache=[tt_kv_cache],
	enable_trace=False,
	start_pos=[start_pos],
	)
	ref_output_slice = ref_output[:, last_token_idx : last_token_idx + 1, :]

	passing, pcc_message = comp_pcc(ref_output_slice, tt_output_torch, pcc)

	logger.info(comp_allclose(ref_output_slice, tt_output_torch))
	logger.info(
	f"passing: {passing}, PCC: {pcc_message} (for last_token_idx: {last_token_idx}, start_pos: {start_pos})"
	)
	assert passing