Safety Attention Kernels

struct AttentionLaunch

Fixed launch dimensions validated against the plugin descriptors before execution.

Public Members

int32_t tokens = {}

int32_t qHeads = {}

int32_t kvHeads = {}

int32_t headSize = {}

int32_t capacity = {}

int32_t ropeTokens = {}

int32_t rotaryDim = {}

int32_t sequences = {}

float scale = {}

uint64_t trt_edgellm::safety::attentionWorkspaceBytes(

AttentionLaunch const &launch

) noexcept

Workspace includes validation/cumulative lengths followed by aligned FP16 rotated Q storage.

cudaError_t trt_edgellm::safety::prepareAttentionKernels() noexcept

Initialization only: prepare statically compiled helpers after explicit FMHA module preparation.

cudaError_t trt_edgellm::safety::launchSafetyLlmAttention(

SafetyAttentionModules &modules,

fmha::PreparedBindings &bindings,

AttentionLaunch const &launchShape,

half const *qkv,

half const *past,

int32_t const *lengths,

float const *rope,

int32_t const *starts,

half *output,

half *present,

void *workspace,

uint64_t *errorMask,

cudaStream_t stream

) noexcept

Execute batch-one FP16 causal attention with contiguous KV, without allocating resources. errorMask is the GPU-accessible bitMask of LWE’s mapped pinned error storage. Host inspection uses its host address after GPU completion.

cudaError_t trt_edgellm::safety::launchSafetyVisualAttention(

SafetyAttentionModules &modules,

fmha::PreparedBindings &bindings,

AttentionLaunch const &launch,

half const *q,

half const *k,

half const *v,

int32_t const *sequenceEnds,

half *output,

void *workspace,

uint64_t *errorMask,

cudaStream_t stream

) noexcept

Execute FP16 bidirectional ragged attention with no KV cache.