File size: 725 Bytes
8acf152 f475609 8acf152 |
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 |
import torch
from ._ops import ops
def get_mla_metadata(seqlens_k: torch.Tensor, s_q: int, h_kv: int):
return ops.get_mla_metadata(seqlens_k, s_q, h_kv)
def mha_fwd_kvcache_mla(
q: torch.Tensor,
kcache: torch.Tensor,
vcache_: torch.Tensor,
head_size_v: int,
seqlens_k: torch.Tensor,
block_table: torch.Tensor,
softmax_scale: float,
is_causal_: bool,
tile_scheduler_metadata: torch.Tensor,
num_splits: torch.Tensor,
) -> torch.Tensor:
return ops.mha_fwd_kvcache_mla(
q,
kcache,
vcache_,
head_size_v,
seqlens_k,
block_table,
softmax_scale,
is_causal_,
tile_scheduler_metadata,
num_splits
)
|