def flash_mla_with_kvcache()

in flash_mla/flash_mla_interface.py [0:0]

• 10 lines of code
• 1 McCabe index (conditional complexity)

def flash_mla_with_kvcache(
    q: torch.Tensor,
    k_cache: torch.Tensor,
    block_table: torch.Tensor,
    cache_seqlens: torch.Tensor,
    head_dim_v: int,
    tile_scheduler_metadata: torch.Tensor,
    num_splits: torch.Tensor,
    softmax_scale: Optional[float] = None,
    causal: bool = False,