Skip to main content

flash_decode_admissible

Function flash_decode_admissible 

Source
pub fn flash_decode_admissible(
    q_shape: &Shape,
    dtype: DType,
    window_size_left: Option<usize>,
    window_size_right: Option<usize>,
    softcap: Option<f32>,
    has_alibi: bool,
    cap: FlashArmCapability,
) -> bool
Expand description

Is a decode-attention region admissible for the CUDA flash arm?

Rejects (→ false, decomposed-only) anything the capacity-K CUDA binding cannot do, matching cost::cost_flash_decoding_cuda’s infeasible set and the kernel’s _can_implement gate:

  • capability unavailable (no CUDA flash kernel / no CUDA device);
  • dtype not f16 / bf16;
  • q not rank-4, seq_q != 1, or head_dim 0 or > 128;
  • sliding window (window_size_{left,right}), softcap, or ALiBi set.

causal is intentionally NOT a disqualifier: with seq_q == 1 the single query attends the whole [0, k_len) prefix, so the causal mask is a no-op the kernel omits.