Class TorchNativeAttentionKernel

java.lang.Object

smile.llm.attention.TorchNativeAttentionKernel

All Implemented Interfaces:

AttentionKernel

public final class TorchNativeAttentionKernel
extends Object
implements AttentionKernel

LibTorch scaled_dot_product_attention backend.

Constructor Summary

Constructors

Constructor	Description
TorchNativeAttentionKernel()	Creates a LibTorch SDPA 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

TorchNativeAttentionKernel

public TorchNativeAttentionKernel()

Creates a LibTorch SDPA 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.