Interface AttentionKernel

All Known Implementing Classes:

FlashInferAttentionKernel, TorchNativeAttentionKernel

public interface AttentionKernel

Attention compute kernel (contiguous SDPA or paged FlashInfer).

Method Summary

Modifier and Type	Method	Description
AttentionBackend	backend()	Returns the backend this kernel implements.
Tensor	forward(Tensor query, Tensor key, Tensor value, Tensor mask, AttentionContext ctx)	Runs attention.

Method Details

forward

Tensor forward(Tensor query,
 Tensor key,
 Tensor value,
 Tensor mask,
 AttentionContext ctx)

Runs attention.

For AttentionBackend.TORCH_NATIVE, query/key/value are contiguous [B,H,S,D] / [B,H,L,D] and mask is optional.

For AttentionBackend.FLASHINFER, query is [B,H,S,D] (or [B,S,H,D] — see kernel docs); key/value may be null when AttentionContext.isPaged() (KV read from the pool).

Parameters:

query - query tensor.

key - key tensor, or null when paged.

value - value tensor, or null when paged.

mask - attention mask, or null.

ctx - scale / causal / paged metadata.

Returns:

attention output matching query layout expected by the caller ([B,H,S,D] for both backends in v1).

backend

AttentionBackend backend()

Returns the backend this kernel implements.

Returns:

backend this kernel implements.