Class FlashInferAttentionKernel

java.lang.Object

smile.llm.attention.FlashInferAttentionKernel

All Implemented Interfaces:

AttentionKernel

public final class FlashInferAttentionKernel
extends Object
implements AttentionKernel

FlashInfer paged BatchPrefill / BatchDecode backend.

Decode (S==1) uses FlashInfer BatchDecodeWithPagedKVCache (vendored headers). Prefill uses GPU page gather + LibTorch SDPA with the same mask semantics as torch_native.

Ragged contiguous prefill (vision tower) uses BatchPrefillWithRaggedKVCache when supported, otherwise batched SDPA inside the FlashInfer module.

Expects AttentionContext.isPaged() with CSR metadata and a FlashInferWorkspace, or AttentionContext.isRaggedContiguous() with cumulative segment indptr. Query layout [B, H, S, D] for paged; [N, H, D] NHD for ragged.

Constructor Summary

Constructors

Constructor	Description
FlashInferAttentionKernel()	Creates a FlashInfer attention kernel.

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.

Methods inherited from class Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, toString, wait, wait, wait

Constructor Details

FlashInferAttentionKernel

public FlashInferAttentionKernel()

Creates a FlashInfer attention kernel.

Method Details

forward

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

Description copied from interface: AttentionKernel

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).

Specified by:

forward in interface AttentionKernel

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

public AttentionBackend backend()

Description copied from interface: AttentionKernel

Returns the backend this kernel implements.

Specified by:

backend in interface AttentionKernel

Returns:

backend this kernel implements.