Gdn Tree Chunk Kernels#
-
bool trt_edgellm::kernel::gdnTreeChunkVerifyEnabled(int32_t treeSize)#
Single source of truth for “is the stateless chunk-form verify used for a
tree of this size?” The plugin verify dispatch and the decoder commit path MUST both gate on this: if they disagreed, one side could skip replay while the other expects replay stash, silently corrupting the committed recurrent state. Callers add their own context guards (plugin: DDTree verify phase; decoder: hybrid-state commit).
- cudaError_t trt_edgellm::kernel::gdnTreeBuildAncestorMasks(
- int32_t const *parentIds,
- uint32_t *masksOut,
- int32_t batch,
- int32_t numNodes,
- int32_t maxDepth,
- cudaStream_t stream
Build packed inclusive ancestor masks from DDTree parent ids. parentIds [GPU int32]: [batch, numNodes], root/padding = -1 (only node 0 may be a root; padding nodes have parent -1 at index > 0 and mask 0). masksOut [GPU uint32]: [batch, numNodes, kGDN_TREE_CHUNK_MASK_WORDS]. Returns the kernel launch status (cudaSuccess on success); the caller must propagate a failure rather than continue with an unwritten output.
- cudaError_t trt_edgellm::kernel::gdnTreeVerifyChunk(
- float const *h0,
- __half const *q,
- __half const *k,
- __half const *v,
- __half const *a,
- __half const *b,
- float const *A_log,
- __half const *dt_bias,
- uint32_t const *masks,
- __half *o,
- void *stash,
- size_t stashBatchStrideBytes,
- int32_t batch,
- int32_t numNodes,
- int32_t h,
- int32_t hv,
- float scale,
- bool useQKL2Norm,
- cudaStream_t stream
Chunk-form tree verify. Reads h0 (fp32, READ-ONLY), emits per-node outputs o (fp16) and the replay stash. No state is written anywhere. Shapes follow the GDN plugin contract: q,k [batch, N, h, 128] fp16; v [batch, N, hv, 128] fp16; a,b [batch, N, hv] fp16; A_log [hv] fp32; dt_bias [hv] fp16; h0 [batch, hv, 128, 128] fp32; o [batch, N, hv, 128] fp16; masks [batch, N, kGDN_TREE_CHUNK_MASK_WORDS] uint32 (inclusive); stash: base pointer of this layer’s intermediate buffer; stride stashBatchStrideBytes between batch rows; nodes packed at gdnTreeStashNodeBytes() intervals. The unused row tail past the last stash cell also holds the verify scratch (KS/QS + prep blocks: align256(MAX_NODES * nodeBytes) + hv*2*MAX_NODES*128 f32 + hv*kGDN_TREE_CHUNK_PREP_WORDS f32); the launcher validates capacity and refuses to launch on overflow. Returns the launch status: cudaSuccess, or the launch/attribute error, or cudaErrorInvalidValue if the scratch does not fit the row stride. The caller MUST propagate a non-success return — continuing would consume garbage verify output and, via the unwritten stash, corrupt the persistent recurrent state on the subsequent replay commit.
- cudaError_t trt_edgellm::kernel::gdnTreeReplayCommitBatched(
- MtpLayerInfo const *deviceLayerInfos,
- int32_t numLayers,
- size_t stashBatchStrideBytes,
- int32_t const *acceptedIndices,
- int32_t const *acceptLengths,
- int32_t batch,
- int32_t maxAcceptLen,
- int32_t numNodes,
- int32_t h,
- int32_t hv,
- cudaStream_t stream
Batched replay commit across ALL recurrent layers in one launch (grid = [batch*hv, numLayers]). Per layer, advances the persistent state (MtpLayerInfo::recurrentDst) along the accepted path using the stash the chunk-form verify wrote at the head of MtpLayerInfo::recurrentSrc. Identical math/op-order to a sequential scan over the accepted nodes. acceptedIndices [GPU int32]: [batch, maxAcceptLen] verify-node indices. acceptLengths [GPU int32]: [batch]. Returns the kernel launch status (cudaSuccess on success); the caller must propagate a failure rather than continue with a partially-committed state.