Dspark Kernels#
- int32_t trt_edgellm::kernel::dsparkMarkovPartialCount(
- int32_t vocabSize
Number of vocab blocks the Markov kernels partition the vocabulary into.
- void trt_edgellm::kernel::dsparkBuildVerifyTokens(
- rt::Tensor const &lastAcceptedTokens,
- rt::Tensor const &draftTokenIds,
- rt::Tensor &verifyTokenIds,
- int32_t batchSize,
- int32_t draftStride,
- int32_t verifyProposalLen,
- cudaStream_t stream
Build DSpark base-verify input IDs.
verify[b, 0] = lastAcceptedTokens[b] verify[b, j + 1] = draftTokenIds[b, j], j in [0, verifyProposalLen)
- void trt_edgellm::kernel::dsparkGreedyAccept(
- rt::Tensor const &baseLogits,
- rt::Tensor const &draftTokenIds,
- rt::Tensor const &proposalLengths,
- rt::Tensor &acceptedTokenIds,
- rt::Tensor &acceptLength,
- rt::Tensor &argmaxScratch,
- int32_t batchSize,
- int32_t draftStride,
- int32_t verifyProposalLen,
- int32_t vocabSize,
- cudaStream_t stream
Greedy DSpark accept for a linear proposal block.
This deterministic path accepts a draft token while it matches the base argmax at the previous verify position, then appends the base bonus token. It is the top-1 analogue of DeepSpec’s rejection-sampling verifier.
- void trt_edgellm::kernel::dsparkFillProposalLengths( )#
Fill logical proposal lengths with a fixed full-block value.
- void trt_edgellm::kernel::dsparkComputeConfidenceAndProposalLengths(
- rt::Tensor const &draftHiddenStates,
- rt::Tensor const &markovW1,
- rt::Tensor const &confidenceWeight,
- rt::Tensor const &confidenceBias,
- rt::Tensor const &firstPrevTokens,
- rt::Tensor const &draftTokenIds,
- rt::Tensor &confidenceScores,
- rt::Tensor &proposalLengths,
- int32_t batchSize,
- int32_t proposalLen,
- int32_t hiddenSize,
- int32_t markovRank,
- bool confidenceWithMarkov,
- float threshold,
- int32_t minProposalLen,
- int32_t maxProposalLen,
- cudaStream_t stream,
- int32_t hiddenStride = 0,
- int32_t hiddenOffset = 0
Compute DSpark confidence scores and threshold-scheduled proposal lengths.
- void trt_edgellm::kernel::dsparkComputeConfidenceAndSPSProposalLengths(
- rt::Tensor const &draftHiddenStates,
- rt::Tensor const &markovW1,
- rt::Tensor const &confidenceWeight,
- rt::Tensor const &confidenceBias,
- rt::Tensor const &firstPrevTokens,
- rt::Tensor const &draftTokenIds,
- rt::Tensor &confidenceScores,
- rt::Tensor &proposalLengths,
- int32_t batchSize,
- int32_t proposalLen,
- int32_t hiddenSize,
- int32_t markovRank,
- bool confidenceWithMarkov,
- float survivalFloor,
- int32_t minProposalLen,
- int32_t maxProposalLen,
- cudaStream_t stream,
- int32_t hiddenStride = 0,
- int32_t hiddenOffset = 0
Compute DSpark confidence scores and SPS-scheduled proposal lengths.
SPS uses draft-side confidence only. It chooses the logical prefix length that maximizes expected output tokens per estimated B100 verify-bucket cost.
- void trt_edgellm::kernel::dsparkComputeConfidenceScores(
- rt::Tensor const &draftHiddenStates,
- rt::Tensor const &markovW1,
- rt::Tensor const &confidenceWeight,
- rt::Tensor const &confidenceBias,
- rt::Tensor const &firstPrevTokens,
- rt::Tensor const &draftTokenIds,
- rt::Tensor &confidenceScores,
- int32_t batchSize,
- int32_t proposalLen,
- int32_t hiddenSize,
- int32_t markovRank,
- bool confidenceWithMarkov,
- cudaStream_t stream,
- int32_t hiddenStride = 0,
- int32_t hiddenOffset = 0
Compute DSpark per-step acceptance confidence scores only (no scheduling).
Retained for chain-side confidence analysis; confidenceScores is [batch, proposalLen].
- void trt_edgellm::kernel::dsparkLogitsToProbabilities(
- rt::Tensor const &logits,
- rt::Tensor &probabilities,
- int32_t rows,
- int32_t vocabSize,
- float temperature,
- int32_t topK,
- float topP,
- cudaStream_t stream
Convert logits to the sampling probability distribution used by DSpark.
This materializes a dense probability row after temperature, top-k, and top-p filtering. It is a correctness-first Track A helper for DSpark probabilistic verification; optimized residual paths can avoid keeping all rows later.
- void trt_edgellm::kernel::dsparkFillUniforms(
- rt::Tensor &uniforms,
- int32_t totalElements,
- uint64_t philoxSeed,
- uint64_t philoxOffset,
- cudaStream_t stream
Fill a GPU tensor with deterministic uniform random values in [0, 1).
- void trt_edgellm::kernel::dsparkVanillaMarkovSample(
- rt::Tensor const &backboneLogits,
- rt::Tensor const &markovW1,
- rt::Tensor const &markovW2,
- rt::Tensor const &firstPrevTokens,
- rt::Tensor const &proposalUniforms,
- rt::Tensor &draftTokenIds,
- rt::Tensor &draftProbabilities,
- rt::Tensor &correctedLogitsScratch,
- rt::Tensor &probabilityScratch,
- int32_t batchSize,
- int32_t proposalLen,
- int32_t vocabSize,
- int32_t markovRank,
- float temperature,
- int32_t topK,
- float topP,
- cudaStream_t stream,
- int32_t logitsStride = 0,
- int32_t logitsOffset = 0
Stochastic DSpark vanilla Markov proposal.
For each proposal step this builds the corrected Markov distribution, materializes it in draftProbabilities [B, proposalLen, vocabSize], and samples one draft token with proposalUniforms [B, proposalLen].
- void trt_edgellm::kernel::dsparkBuildMarkovLogits(
- rt::Tensor const &backboneLogits,
- rt::Tensor const &markovW1,
- rt::Tensor const &markovW2,
- rt::Tensor const &firstPrevTokens,
- rt::Tensor const &draftTokenIds,
- rt::Tensor &correctedLogitsScratch,
- int32_t batchSize,
- int32_t step,
- int32_t proposalLen,
- int32_t vocabSize,
- int32_t markovRank,
- cudaStream_t stream,
- int32_t logitsStride = 0,
- int32_t logitsOffset = 0
Build corrected DSpark Markov logits for one proposal step.
- bool trt_edgellm::kernel::dsparkFusedGreedySupported(
- int32_t markovRank
True when the fused greedy Markov step kernel supports this markov rank.
- void trt_edgellm::kernel::dsparkMarkovGreedyFusedStep(
- rt::Tensor const &backboneLogits,
- rt::Tensor const &markovW1,
- rt::Tensor const &markovW2,
- rt::Tensor const &firstPrevTokens,
- rt::Tensor &greedySlots,
- rt::Tensor *stackedLogits,
- int32_t stackedDepthRow,
- int32_t batchSize,
- int32_t step,
- int32_t proposalLen,
- int32_t vocabSize,
- int32_t markovRank,
- cudaStream_t stream,
- int32_t logitsStride = 0,
- int32_t logitsOffset = 0
Fused greedy Markov proposal step: correction + exact top-1 in one launch.
Reads the previous step’s winner from greedySlots (packed orderable-float key, ties resolve to the lowest vocab index), so no host-visible token round trip is needed between steps. greedySlots is [batch, proposalLen] UINT64 and must be zeroed once per round before step 0. When stackedLogits is non-null, the corrected full-vocab row is also written to depth row stackedDepthRow of the [batch, proposalLen + 1, vocab] tree candidate buffer; chain mode passes null and never materializes the row.
- void trt_edgellm::kernel::dsparkQuantizeMarkovW2Fp8(
- rt::Tensor const &markovW2,
- rt::Tensor &w2Fp8,
- rt::Tensor &w2RowScales,
- int32_t vocabSize,
- int32_t markovRank,
- cudaStream_t stream
Quantize markov_w2 [V, R] FP16 to FP8 E4M3 with one FP16 scale per row.
w2Fp8 is [V, R] UINT8 (E4M3 bytes of w / rowScale); w2RowScales is [V] FP16 with rowScale = rowAbsMax / 448 so each row spans the full E4M3 range. Requires R % 16 == 0.
- void trt_edgellm::kernel::dsparkMarkovGreedyFusedStepFp8(
- rt::Tensor const &backboneLogits,
- rt::Tensor const &markovW1,
- rt::Tensor const &w2Fp8,
- rt::Tensor const &w2RowScales,
- rt::Tensor const &firstPrevTokens,
- rt::Tensor &greedySlots,
- rt::Tensor *stackedLogits,
- int32_t stackedDepthRow,
- int32_t batchSize,
- int32_t step,
- int32_t proposalLen,
- int32_t vocabSize,
- int32_t markovRank,
- cudaStream_t stream,
- int32_t logitsStride = 0,
- int32_t logitsOffset = 0
FP8-W2 variant of the fused greedy Markov step (same contract as dsparkMarkovGreedyFusedStep). Two tokens per warp, one 16-byte load per lane, hardware fp8x2 -> half2 conversion. Requires markovRank == 16 * 2^k so the per-token lane groups tile a warp exactly.
- void trt_edgellm::kernel::dsparkFinalizeGreedyDraftTokens( )#
Unpack the per-step greedy winners in greedySlots into draftTokenIds [batch, proposalLen].
- void trt_edgellm::kernel::dsparkSampleProbabilityRows(
- rt::Tensor const &probabilityScratch,
- rt::Tensor const &proposalUniforms,
- rt::Tensor &draftTokenIds,
- int32_t batchSize,
- int32_t step,
- int32_t proposalLen,
- int32_t vocabSize,
- cudaStream_t stream
Sample one token per row from probabilityScratch [B, vocabSize].
- void trt_edgellm::kernel::dsparkStoreDraftStepProbabilities(
- rt::Tensor const &probabilityScratch,
- rt::Tensor &draftProbabilities,
- int32_t batchSize,
- int32_t step,
- int32_t proposalLen,
- int32_t vocabSize,
- cudaStream_t stream
Store one proposal step probability row into draftProbabilities [B, P, V].
- void trt_edgellm::kernel::dsparkNormalizeTopKRows(
- rt::Tensor const &topKValues,
- rt::Tensor &topKProbabilities,
- int32_t rows,
- int32_t topK,
- float temperature,
- cudaStream_t stream
Normalize selected top-k logits into sparse probabilities.
- void trt_edgellm::kernel::dsparkStoreDraftStepTop1(
- rt::Tensor const &top1Indices,
- rt::Tensor &draftTokenIds,
- int32_t batchSize,
- int32_t step,
- int32_t proposalLen,
- cudaStream_t stream
Store one selected top-1 token per batch row into draftTokenIds [B, P].
- void trt_edgellm::kernel::dsparkSampleTopKRowsAndStore(
- rt::Tensor const &topKValues,
- rt::Tensor const &topKIndices,
- rt::Tensor const &proposalUniforms,
- rt::Tensor &draftTokenIds,
- rt::Tensor &draftTopKProbabilities,
- rt::Tensor &draftTopKIndices,
- int32_t batchSize,
- int32_t step,
- int32_t proposalLen,
- int32_t topK,
- float temperature,
- cudaStream_t stream
Sample one DSpark draft token from selected top-k logits and store sparse draft probabilities/indices.
- void trt_edgellm::kernel::dsparkProbabilisticAccept(
- rt::Tensor const &targetProbabilities,
- rt::Tensor const &draftProbabilities,
- rt::Tensor const &draftTokenIds,
- rt::Tensor const &proposalLengths,
- rt::Tensor const &acceptUniforms,
- rt::Tensor &acceptedTokenIds,
- rt::Tensor &acceptLength,
- int32_t batchSize,
- int32_t draftStride,
- int32_t verifyProposalLen,
- int32_t vocabSize,
- cudaStream_t stream
DSpark probabilistic verifier with residual sampling.
targetProbabilities is [B, proposalLen + 1, vocabSize]. draftProbabilities is [B, proposalLen, vocabSize]. acceptUniforms is [B, 2 * proposalLen + 1]: first proposalLen values drive accept/reject, the next proposalLen values drive residual sampling for a rejection at the matching step, and the final value drives the bonus-token sample when the full proposal is accepted.