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
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
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
Compute DSpark per-step acceptance confidence scores only (no scheduling).
Used by DDTree drafting to bias tree growth; 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
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
Build corrected DSpark Markov logits for one proposal step.
- 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::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::dsparkSparseTopKAccept(
- rt::Tensor const &targetTopKProbabilities,
- rt::Tensor const &targetTopKIndices,
- rt::Tensor const &draftTopKProbabilities,
- rt::Tensor const &draftTopKIndices,
- 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 targetTopK,
- int32_t draftTopK,
- cudaStream_t stream
DSpark probabilistic verifier over sparse top-k target/draft supports.
- 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.