in benchmark/bench_flash_mla.py [0:0]
def _mla_attn( q_nope, q_pe, kv_c_cache, k_pe_cache, attn_logits, req_to_tokens, b_seq_len, num_kv_splits, sm_scale, page_size,