Skip to main content

DispatchKey

Struct DispatchKey 

Source
pub struct DispatchKey {
    pub arch: SmArch,
    pub dtype: DType,
    pub head_dim: u32,
    pub causal: bool,
    pub varlen: bool,
    pub sliding_window: Option<u32>,
    pub alibi: bool,
    pub sink: u32,
    pub paged: bool,
    pub gqa_ratio: u32,
}
Expand description

Cell key for the FlashAttention dispatch table.

Every field directly affects the generated CUDA C++ template instantiation — flipping any one of them changes the resulting cubin. The table refuses to resolve unsupported combinations (e.g. fp8 on Sm80, head_dim > 256).

Fields§

§arch: SmArch

Target SM architecture.

§dtype: DType

Element type for Q/K/V.

§head_dim: u32

Per-head dimension (D). Supported: 64, 80, 96, 128, 192, 256.

§causal: bool

Causal masking — autoregressive attention.

§varlen: bool

Variable-length (cu_seqlens). When false, batched attention with uniform seqlen.

§sliding_window: Option<u32>

Sliding-window size; None means full attention. Window size is the number of past tokens each query attends to.

§alibi: bool

ALiBi linear-position biases.

§sink: u32

Number of “sink” tokens (StreamingLLM); each query unconditionally attends to the first sink keys regardless of sliding_window.

§paged: bool

vLLM-style paged KV-cache.

§gqa_ratio: u32

Q heads per KV head. 1 = MHA, >1 = GQA, equal to num_heads = MQA.

Implementations§

Source§

impl DispatchKey

Source

pub fn validate_fwd(&self) -> Result<(), DispatchError>

Validate the cell for a forward path. Returns Err for unreachable combinations.

Source

pub fn validate_bwd(&self) -> Result<(), DispatchError>

Validate the cell for a backward path. Currently the same as forward, but kept distinct so we can refuse e.g. fp8 backward (numerically too lossy in the stock FA3) without affecting the forward whitelist.

Source

pub fn validate_paged(&self) -> Result<(), DispatchError>

Validate the cell for a paged forward path.

Source

pub fn stable_hash(&self) -> u64

Stable 64-bit hash of the key. Useful as a cubin-cache index alongside the kernel-name string.

Source

pub fn kernel_name(&self) -> String

Build the canonical mangled kernel-name expression. Mirrors the FA2/FA3 csrc naming convention so we can resolve it via NVRTC’s nvrtcGetLoweredName.

Trait Implementations§

Source§

impl Clone for DispatchKey

Source§

fn clone(&self) -> DispatchKey

Returns a duplicate of the value. Read more
1.0.0 (const: unstable) · Source§

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source. Read more
Source§

impl Debug for DispatchKey

Source§

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter. Read more
Source§

impl Hash for DispatchKey

Source§

fn hash<__H: Hasher>(&self, state: &mut __H)

Feeds this value into the given Hasher. Read more
1.3.0 · Source§

fn hash_slice<H>(data: &[Self], state: &mut H)
where H: Hasher, Self: Sized,

Feeds a slice of this type into the given Hasher. Read more
Source§

impl PartialEq for DispatchKey

Source§

fn eq(&self, other: &DispatchKey) -> bool

Tests for self and other values to be equal, and is used by ==.
1.0.0 (const: unstable) · Source§

fn ne(&self, other: &Rhs) -> bool

Tests for !=. The default implementation is almost always sufficient, and should not be overridden without very good reason.
Source§

impl Copy for DispatchKey

Source§

impl Eq for DispatchKey

Source§

impl StructuralPartialEq for DispatchKey

Auto Trait Implementations§

Blanket Implementations§

Source§

impl<T> Any for T
where T: 'static + ?Sized,

Source§

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more
Source§

impl<T> Borrow<T> for T
where T: ?Sized,

Source§

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more
Source§

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source§

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more
Source§

impl<T> CloneToUninit for T
where T: Clone,

Source§

unsafe fn clone_to_uninit(&self, dest: *mut u8)

🔬This is a nightly-only experimental API. (clone_to_uninit)
Performs copy-assignment from self to dest. Read more
Source§

impl<Q, K> Equivalent<K> for Q
where Q: Eq + ?Sized, K: Borrow<Q> + ?Sized,

Source§

fn equivalent(&self, key: &K) -> bool

Checks if this value is equivalent to the given key. Read more
Source§

impl<Q, K> Equivalent<K> for Q
where Q: Eq + ?Sized, K: Borrow<Q> + ?Sized,

Source§

fn equivalent(&self, key: &K) -> bool

Compare self to key and return true if they are equal.
Source§

impl<Q, K> Equivalent<K> for Q
where Q: Eq + ?Sized, K: Borrow<Q> + ?Sized,

Source§

fn equivalent(&self, key: &K) -> bool

Checks if this value is equivalent to the given key. Read more
Source§

impl<T> From<T> for T

Source§

fn from(t: T) -> T

Returns the argument unchanged.

Source§

impl<T> Instrument for T

Source§

fn instrument(self, span: Span) -> Instrumented<Self>

Instruments this type with the provided Span, returning an Instrumented wrapper. Read more
Source§

fn in_current_span(self) -> Instrumented<Self>

Instruments this type with the current Span, returning an Instrumented wrapper. Read more
Source§

impl<T, U> Into<U> for T
where U: From<T>,

Source§

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source§

impl<T> ToOwned for T
where T: Clone,

Source§

type Owned = T

The resulting type after obtaining ownership.
Source§

fn to_owned(&self) -> T

Creates owned data from borrowed data, usually by cloning. Read more
Source§

fn clone_into(&self, target: &mut T)

Uses borrowed data to replace owned data, usually by cloning. Read more
Source§

impl<T, U> TryFrom<U> for T
where U: Into<T>,

Source§

type Error = Infallible

The type returned in the event of a conversion error.
Source§

fn try_from(value: U) -> Result<T, <T as TryFrom<U>>::Error>

Performs the conversion.
Source§

impl<T, U> TryInto<U> for T
where U: TryFrom<T>,

Source§

type Error = <U as TryFrom<T>>::Error

The type returned in the event of a conversion error.
Source§

fn try_into(self) -> Result<U, <U as TryFrom<T>>::Error>

Performs the conversion.
Source§

impl<T> WithSubscriber for T

Source§

fn with_subscriber<S>(self, subscriber: S) -> WithDispatch<Self>
where S: Into<Dispatch>,

Attaches the provided Subscriber to this type, returning a WithDispatch wrapper. Read more
Source§

fn with_current_subscriber(self) -> WithDispatch<Self>

Attaches the current default Subscriber to this type, returning a WithDispatch wrapper. Read more
Source§

impl<T> Extension for T
where T: Any + Send + Sync,