Class GlideInput

colossalai/inference/spec/struct.py:33–56

Dataclass for Glide Models (e.g. `colossalai/inference/modeling/models/glide_llama.py`). Used for pack data that will be used during glimpsing KV Caches of the main model. Args: block_tables (torch.Tensor): [num_seqs, max_blocks_per_seq] The block table of KV Caches. large_k

@dataclass
class GlideInput:
    """Dataclass for Glide Models (e.g. `colossalai/inference/modeling/models/glide_llama.py`).
    Used for pack data that will be used during glimpsing KV Caches of the main model.

    Args:
        block_tables (torch.Tensor): [num_seqs, max_blocks_per_seq] The block table of KV Caches.
        large_k_cache (torch.Tensor): [num_blocks, num_kv_heads, block_size, head_size]
            Blocked key cache of the main model
        large_v_cache (torch.Tensor): Blocked value cache of the main model. It has the same shape as k cache.
        sequence_lengths (torch.Tensor): [num_seqs] Sequence lengths of the current batch.
    """

    block_tables: torch.Tensor = None
    large_k_cache: torch.Tensor = None
    large_v_cache: torch.Tensor = None
    sequence_lengths: torch.Tensor = None
    n_spec_tokens: int = 5

    @property
    def glimpse_ready(self):
        return all(
            attr is not None
            for attr in [self.block_tables, self.large_k_cache, self.large_v_cache, self.sequence_lengths]
        )