↓ 1 callersFunctionscaled_dot_product_attention(
batch_idx: int,
query: torch.Tensor, # [h_q, s_q, d]
kv: torch.Tensor, # [h_
tests/test_flash_mla_dense_decoding.py:85
↓ 1 callersFunctiontest_flash_attention(b, mean_sq, mean_sk, varlen, h, h_k, d, dv, causal, window, has_bwd, check_correctness: bool = True)
tests/test_fmha_sm100.py:50
Function_mla_attn_kernel(
Q_nope,
Q_pe,
Kv_c_cache,
K_pe_cache,
Req_to_tokens,
B_seq_len,
O,
sm_scale,
benchmark/bench_flash_mla.py:136
Function_mla_softmax_reducev_kernel(
Logits,
B_seq_len,
O,
stride_l_b,
stride_l_h,
stride_l_s,
stride_o_b,
stride
benchmark/bench_flash_mla.py:274
Functioncheck_is_allclose_comparator(name: str, ans: torch.Tensor, ref: torch.Tensor, out: torch.Tensor, abs_tol: float = 1e-5, rel_tol: float = 1
tests/kernelkit/compare.py:94
Functionflash_mla_sparse_fwd Sparse attention prefill kernel Args: q: [s_q, h_q, d_qk], bfloat16 kv: [s_kv, h_kv, d_qk], bfloat16 indices: [s_q,
flash_mla/flash_mla_interface.py:176