File size: 1,501 Bytes
d91766b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | import torch
from dataclasses import dataclass
from diffulex.strategy_template.token_merging_multi_block.attention.metadata import (
TokenMergingMultiBlockAttnMetaDataTemplate,
)
@dataclass
class DMaxAttnMetaData(TokenMergingMultiBlockAttnMetaDataTemplate):
def __post_init__(self):
self.init_multi_block()
self.reset_token_merging()
DMAX_ATTN_METADATA = DMaxAttnMetaData()
def fetch_dmax_attn_metadata() -> DMaxAttnMetaData:
return DMAX_ATTN_METADATA
def set_dmax_attn_metadata(
is_prefill: bool = False,
cu_seqlens_q: torch.Tensor | None = None,
cu_seqlens_k: torch.Tensor | None = None,
max_seqlen_q: int = 0,
max_seqlen_k: int = 0,
slot_mapping: torch.Tensor | None = None,
context_lens: torch.Tensor | None = None,
page_tables: torch.Tensor | None = None,
page_size: int = 32,
block_size: int = 32,
kv_cache_layout: str = "unified",
) -> None:
global DMAX_ATTN_METADATA
DMAX_ATTN_METADATA = DMaxAttnMetaData(
is_prefill=is_prefill,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
max_seqlen_q=max_seqlen_q,
max_seqlen_k=max_seqlen_k,
slot_mapping=slot_mapping,
context_lens=context_lens,
page_tables=page_tables,
page_size=page_size,
block_size=block_size,
kv_cache_layout=kv_cache_layout,
)
def reset_dmax_attn_metadata() -> None:
global DMAX_ATTN_METADATA
DMAX_ATTN_METADATA = DMaxAttnMetaData()
|