diff --git a/cpp/tensorrt_llm/kernels/IndexerTopK.h b/cpp/tensorrt_llm/kernels/IndexerTopK.h index 6e6e7b29b059..025c7a19b404 100644 --- a/cpp/tensorrt_llm/kernels/IndexerTopK.h +++ b/cpp/tensorrt_llm/kernels/IndexerTopK.h @@ -40,11 +40,18 @@ int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitW /// - Insertion sort (N < kSortingAlgorithmThreshold) /// - Radix sort (kSortingAlgorithmThreshold ≤ N < splitWork) /// - Radix split-work (N ≥ splitWork — uses outLogitsAux / outIndicesAux) +/// +/// @param rowKvLens optional [numRows] per-row causal KV extent, in uncompressed +/// token space. Non-null selects the ragged layout: rows no longer map to +/// requests by `rowIdx / next_n`, so the extent every row may attend to is +/// supplied directly instead of being reconstructed from `next_n`. Null keeps +/// the uniform arithmetic, bit-identical to before. Mutually exclusive with +/// `preIdx` (GVR's hint is request-indexed through `next_n` as well). void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, float* heuristicScratch = nullptr, int const compressRatio = 1, - cudaStream_t const stream = 0); + int const* rowKvLens = nullptr, cudaStream_t const stream = 0); /// bf16 indexer TopK decode — same dispatch axes as the fp32 entry, except /// kBsL2 uses sizeof(__nv_bfloat16) bytes/elem (L2 footprint is half) and @@ -59,13 +66,14 @@ void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, in int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, int const compressRatio = 1, - cudaStream_t const stream = 0); + int const* rowKvLens = nullptr, cudaStream_t const stream = 0); /// fp16 indexer TopK decode — see bf16 overload for dispatcher contract. void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, - __half* heuristicScratch = nullptr, int const compressRatio = 1, cudaStream_t const stream = 0); + __half* heuristicScratch = nullptr, int const compressRatio = 1, int const* rowKvLens = nullptr, + cudaStream_t const stream = 0); void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, int const numRows, int const numColumns, int const stride0, int const stride1, int const topK = 2048, diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu index 1457efeb3138..361df4bacaa2 100644 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu @@ -297,7 +297,7 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv, int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens, int32_t const* __restrict__ cu_seq_lens, int32_t const* __restrict__ cu_kv_comp, int page_size, int max_blocks, - int out_elem_bytes) + int out_elem_bytes, int32_t const* __restrict__ new_tokens_per_seq = nullptr) { using KvScoreElemT = typename std::conditional::type; using StateElemT = typename std::conditional::type; @@ -337,7 +337,16 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo int const eff_tid = head_blk * NTHRD_INNER + tid; int const kv_len = kv_lens[batch_idx]; - int const sp = kv_len - NEXT_N; + // How many tokens this request appended this step. Uniform batches append + // NEXT_N each, which is what the template bound encodes. Under DSpark's + // ragged verification each request appends its own window instead, so the + // count is read per request and NEXT_N degrades to a compile-time *upper* + // bound -- which is all the loops below need it to be, because both are + // already guarded (Phase 1 by token_idx < kv_len, Phase 3 by + // num_compressions). Keeping NEXT_N as the bound is what preserves the + // full unroll on the uniform path. + int const nn = (new_tokens_per_seq != nullptr) ? new_tokens_per_seq[batch_idx] : NEXT_N; + int const sp = kv_len - nn; int const in_off = cu_seq_lens[batch_idx]; int const out_off = cu_kv_comp[batch_idx]; int64_t const page_sd = static_cast(page_size) * STATE_DIM; @@ -413,7 +422,7 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo // ================================================================ // Phase 2: Count how many complete compression windows finished. // ================================================================ - int last_token_idx = sp + NEXT_N - 1; + int last_token_idx = sp + nn - 1; int num_compressions = (last_token_idx + 1) / COMPRESS_RATIO - sp / COMPRESS_RATIO; // ================================================================ @@ -663,7 +672,8 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo // Generate explicit template instantiations. #define INST_DECODE(HD, KV_EB, STATE_EB, CR, NN, NRW) \ template __global__ void pagedKvCompressKernel(void const*, float const*, void*, \ - void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int, int, int); + void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int, int, int, \ + int32_t const*); FOREACH_DECODE_CONFIG(INST_DECODE) #undef INST_DECODE @@ -680,6 +690,7 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens, int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, int max_blocks, int head_dim, int compress_ratio, int next_n, int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes, + int32_t const* new_tokens_per_seq, cudaStream_t stream) { TLLM_CHECK_WITH_INFO( @@ -731,7 +742,7 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k { \ pagedKvCompressKernel<<>>(kv_score, ape, \ paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp, \ - page_size, max_blocks, out_elem_bytes); \ + page_size, max_blocks, out_elem_bytes, new_tokens_per_seq); \ return; \ } FOREACH_DECODE_CONFIG(TRY_LAUNCH) diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h index 0d83628e1b9b..df5d2826b03b 100644 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h @@ -47,6 +47,10 @@ void pagedKvCompressLaunch(void const* kv_score, // [m, 2*state_dim] (bf16 or f int kv_score_elem_bytes, // bytes per element for kv_score (2=bf16, 4=fp32) int state_elem_bytes, // bytes per element for paged state (2=bf16, 4=fp32) int out_elem_bytes, // bytes per element for output + int32_t const* new_tokens_per_seq, // [bsz] or nullptr; per-request tokens appended this + // step. nullptr means every request appended next_n + // (the uniform case); non-null makes next_n an upper + // bound only, for DSpark ragged verification. cudaStream_t stream); // Prefill kernel: bulk compression with per-token gather/scatter + state update. diff --git a/cpp/tensorrt_llm/kernels/indexerTopK.cu b/cpp/tensorrt_llm/kernels/indexerTopK.cu index d63f203a048c..12c77af6664c 100644 --- a/cpp/tensorrt_llm/kernels/indexerTopK.cu +++ b/cpp/tensorrt_llm/kernels/indexerTopK.cu @@ -654,7 +654,8 @@ template static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(InputT const* logits, int const* seqLens, int* outIndices, int stride0, int stride1, int const topK, int next_n, int compressRatio, - float* outLogits = nullptr, int const numBlocksToMerge = 0, int const* indices = nullptr) + float* outLogits = nullptr, int const numBlocksToMerge = 0, int const* indices = nullptr, + int const* rowKvLens = nullptr) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); @@ -664,8 +665,27 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I // The range of logits within the row. int rowStart = 0; - int seq_len = seqLens[rowIdx / next_n]; - int actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + // Two ways to learn how far back this query row may attend. + // + // Uniform (rowKvLens == nullptr): every request contributes exactly next_n + // rows, so the request index and the row's offset inside its verify window + // both fall out of integer division, and the causal extent is the request's + // sequence length walked back to this offset. + // + // Ragged (rowKvLens != nullptr): requests contribute different numbers of + // rows, so neither rowIdx / next_n nor rowIdx % next_n means anything. The + // caller precomputes the extent per row -- it is the only quantity this + // kernel ever derived from next_n, which is why one array is enough. + int actual_kv_len; + if (rowKvLens != nullptr) + { + actual_kv_len = rowKvLens[rowIdx]; + } + else + { + int seq_len = seqLens[rowIdx / next_n]; + actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + } int rowEnd = actual_kv_len / compressRatio; // Local pointers to this block @@ -890,9 +910,16 @@ int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitW void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, - int const preIdxCount, float* heuristicScratch, int const compressRatio, cudaStream_t const stream) + int const preIdxCount, float* heuristicScratch, int const compressRatio, int const* rowKvLens, + cudaStream_t const stream) { constexpr int kNumThreadsPerBlock = 512; + // GVR consumes previous-step Top-K hints indexed by request, which it also + // recovers from next_n. Ragged batches therefore cannot take that path; the + // caller is expected to have disabled the hint, and this makes the + // requirement explicit rather than letting a stale preIdx through. + TLLM_CHECK_WITH_INFO(rowKvLens == nullptr || preIdx == nullptr, + "indexer TopK decode: ragged rowKvLens is incompatible with the GVR preIdx hint"); int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; // ======================================================================== @@ -1016,7 +1043,7 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic config.attrs = attrs; cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); + compressRatio, nullptr, 0, nullptr, rowKvLens); } else { @@ -1031,7 +1058,7 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic config_part1.attrs = attrs; cudaLaunchKernelEx(&config_part1, kernel_instance_part1, logits, seqLens, outIndicesAux, stride0, stride1, topK, - next_n, compressRatio, outLogitsAux, 0, nullptr); + next_n, compressRatio, outLogitsAux, 0, nullptr, rowKvLens); constexpr int kNumThreadsPerBlockMerge = 1024; auto* kernel_instance_part2 = &topKPerRowDecode; @@ -1043,8 +1070,12 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic config_part2.numAttrs = 1; config_part2.attrs = attrs; + // The merge pass overrides rowEnd with numBlocksToMerge * topK, so the + // extent is unused -- but rowKvLens still has to be forwarded, because + // otherwise the uniform branch would index seqLens by rowIdx / next_n, + // and on a ragged batch next_n is not a divisor of the row count. cudaLaunchKernelEx(&config_part2, kernel_instance_part2, outLogitsAux, seqLens, indices, blocksPerRow * topK, 1, - topK, next_n, 1, nullptr, blocksPerRow, outIndicesAux); + topK, next_n, 1, nullptr, blocksPerRow, outIndicesAux, rowKvLens); } sync_check_cuda_error(stream); } @@ -1079,10 +1110,14 @@ template void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, int const preIdxCount, InputT* heuristicScratch, int const compressRatio, - cudaStream_t const stream) + int const* rowKvLens, cudaStream_t const stream) { static_assert(std::is_same_v || std::is_same_v, "invokeIndexerTopKDecodeDtype is for bf16/fp16 only"); + // See the fp32 entry: GVR's hint is request-indexed via next_n, so it and a + // ragged row map cannot both be live. + TLLM_CHECK_WITH_INFO(rowKvLens == nullptr || preIdx == nullptr, + "indexer TopK decode: ragged rowKvLens is incompatible with the GVR preIdx hint"); constexpr int kNumThreadsPerBlock = 512; int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; @@ -1123,7 +1158,7 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* config.attrs = attrs; cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); + compressRatio, nullptr, 0, nullptr, rowKvLens); } else if (numColumns < effectiveSplitWorkThreshold) { @@ -1143,7 +1178,7 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* config.attrs = attrs; cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); + compressRatio, nullptr, 0, nullptr, rowKvLens); } else { @@ -1162,19 +1197,20 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, int const preIdxCount, - __nv_bfloat16* heuristicScratch, int const compressRatio, cudaStream_t const stream) + __nv_bfloat16* heuristicScratch, int const compressRatio, int const* rowKvLens, cudaStream_t const stream) { invokeIndexerTopKDecodeDtype<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, - stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); + stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, rowKvLens, + stream); } void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, int const preIdxCount, __half* heuristicScratch, int const compressRatio, - cudaStream_t const stream) + int const* rowKvLens, cudaStream_t const stream) { invokeIndexerTopKDecodeDtype<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); + stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, rowKvLens, stream); } void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, diff --git a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp index fc29041bbf93..60ca03a75e05 100644 --- a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp @@ -38,7 +38,8 @@ namespace torch_ext void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, th::Tensor const& indices, int64_t next_n, int64_t index_topk, std::optional const& pre_idx, std::optional const& heuristic_scratch, int64_t compress_ratio, - std::optional const& radix_aux_indices, std::optional const& radix_aux_logits) + std::optional const& radix_aux_indices, std::optional const& radix_aux_logits, + std::optional const& row_kv_lens) { TORCH_CHECK(compress_ratio > 0, "compress_ratio must be greater than 0"); @@ -53,8 +54,17 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t auto const inputSize = logits.sizes(); auto const numRows64 = inputSize[0]; auto const numColumns64 = inputSize[1]; - TORCH_CHECK( - seq_lens.size(0) * next_n == numRows64, "seq_lens length multiplied by next_n must equal logits.size(0)"); + // Uniform batches must tile exactly: every request owns next_n rows. A + // ragged batch cannot satisfy that by construction, and -- crucially -- can + // still satisfy it *by accident* whenever the padded row count happens to be + // divisible by the row count, at which point the kernel would silently + // attribute rows to the wrong requests. So the ragged layout replaces the + // check rather than loosening it. + if (!row_kv_lens.has_value()) + { + TORCH_CHECK( + seq_lens.size(0) * next_n == numRows64, "seq_lens length multiplied by next_n must equal logits.size(0)"); + } TORCH_CHECK(indices.size(0) == numRows64, "indices first dimension must match logits.size(0)"); TORCH_CHECK(indices.size(1) >= index_topk, "indices second dimension must be at least index_topk"); TORCH_CHECK(seq_lens.is_contiguous(), "seq_lens must be contiguous"); @@ -70,6 +80,22 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t TORCH_CHECK(logits_stride_0 >= 0, "logits_stride_0 must be greater than or equal to 0"); TORCH_CHECK(logits_stride_1 >= 0, "logits_stride_1 must be greater than or equal to 0"); + int32_t const* rowKvLensPtr = nullptr; + if (row_kv_lens.has_value()) + { + auto const& rowKvLensTensor = row_kv_lens.value(); + TORCH_CHECK(rowKvLensTensor.is_cuda(), "row_kv_lens must be a CUDA tensor"); + TORCH_CHECK( + rowKvLensTensor.device() == logits.device(), "row_kv_lens must be on the same device as logits"); + TORCH_CHECK(rowKvLensTensor.is_contiguous(), "row_kv_lens must be contiguous"); + TORCH_CHECK(rowKvLensTensor.dim() == 1, "row_kv_lens must be a 1D Tensor"); + TORCH_CHECK(rowKvLensTensor.scalar_type() == at::ScalarType::Int, "row_kv_lens must be int32"); + TORCH_CHECK(rowKvLensTensor.size(0) == numRows64, + "row_kv_lens must have one entry per logits row (got ", rowKvLensTensor.size(0), " for ", numRows64, + " rows)"); + rowKvLensPtr = rowKvLensTensor.data_ptr(); + } + int32_t const* preIdxPtr = nullptr; int32_t preIdxStride = 0; int32_t preIdxCount = 0; @@ -162,7 +188,8 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t tk::invokeIndexerTopKDecode(logits.data_ptr(), seq_lens.data_ptr(), indices.data_ptr(), aux_logits_ptr, aux_indices_ptr, splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, preIdxStride, - preIdxCount, static_cast(heuristicScratchPtr), static_cast(compress_ratio), stream); + preIdxCount, static_cast(heuristicScratchPtr), static_cast(compress_ratio), rowKvLensPtr, + stream); } else if (logits_dtype == at::ScalarType::BFloat16) { @@ -170,14 +197,14 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t seq_lens.data_ptr(), indices.data_ptr(), splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, preIdxStride, preIdxCount, static_cast<__nv_bfloat16*>(heuristicScratchPtr), - static_cast(compress_ratio), stream); + static_cast(compress_ratio), rowKvLensPtr, stream); } else // Half { tk::invokeIndexerTopKDecode(reinterpret_cast<__half const*>(logits.data_ptr()), seq_lens.data_ptr(), indices.data_ptr(), splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, preIdxStride, preIdxCount, - static_cast<__half*>(heuristicScratchPtr), static_cast(compress_ratio), stream); + static_cast<__half*>(heuristicScratchPtr), static_cast(compress_ratio), rowKvLensPtr, stream); } } @@ -228,7 +255,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def( "indexer_topk_decode(Tensor logits, Tensor seq_lens, Tensor indices, int next_n, int index_topk=2048, " "Tensor? pre_idx=None, Tensor? heuristic_scratch=None, int compress_ratio=1, " - "Tensor? radix_aux_indices=None, Tensor? radix_aux_logits=None) -> ()"); + "Tensor? radix_aux_indices=None, Tensor? radix_aux_logits=None, Tensor? row_kv_lens=None) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/cpp/tensorrt_llm/thop/compressorOp.cpp b/cpp/tensorrt_llm/thop/compressorOp.cpp index 5fdfb5a85806..2b56924df0be 100644 --- a/cpp/tensorrt_llm/thop/compressorOp.cpp +++ b/cpp/tensorrt_llm/thop/compressorOp.cpp @@ -36,12 +36,27 @@ void compressorPagedKvCompressOp(torch::Tensor kv_score, // [m, 2*state_dim] bf1 torch::Tensor kv_lens, // [bsz] int32 torch::Tensor cu_seq_lens, // [bsz+1] int32 torch::Tensor cu_kv_comp, // [bsz+1] int32 - int64_t batch_size, int64_t page_size, int64_t head_dim, int64_t compress_ratio, int64_t next_n) + int64_t batch_size, int64_t page_size, int64_t head_dim, int64_t compress_ratio, int64_t next_n, + std::optional const& new_tokens_per_seq) { constexpr int64_t kMinNextN = 1; constexpr int64_t kMaxNextN = 8; TORCH_CHECK(next_n >= kMinNextN && next_n <= kMaxNextN, "next_n must be in [1, 8], got ", next_n); + // Ragged batches pass the batch maximum as next_n and the real per-request + // counts here; the kernel then treats next_n purely as a loop bound. + int32_t const* new_tokens_ptr = nullptr; + if (new_tokens_per_seq.has_value()) + { + auto const& t = new_tokens_per_seq.value(); + TORCH_CHECK(t.is_cuda(), "new_tokens_per_seq must be a CUDA tensor"); + TORCH_CHECK(t.is_contiguous(), "new_tokens_per_seq must be contiguous"); + TORCH_CHECK(t.scalar_type() == torch::kInt32, "new_tokens_per_seq must be int32"); + TORCH_CHECK(t.dim() == 1 && t.size(0) >= batch_size, + "new_tokens_per_seq must be 1D with at least batch_size entries"); + new_tokens_ptr = t.data_ptr(); + } + auto stream = at::cuda::getCurrentCUDAStream(); int kv_score_eb = static_cast(kv_score.element_size()); int state_eb = static_cast(paged_kv.element_size()); @@ -54,7 +69,7 @@ void compressorPagedKvCompressOp(torch::Tensor kv_score, // [m, 2*state_dim] bf1 kv_lens.data_ptr(), cu_seq_lens.data_ptr(), cu_kv_comp.data_ptr(), static_cast(batch_size), static_cast(page_size), static_cast(block_table_kv.size(1)), static_cast(head_dim), static_cast(compress_ratio), static_cast(next_n), kv_score_eb, state_eb, - out_eb, stream); + out_eb, new_tokens_ptr, stream); } // Prefill kernel: bulk compression with state update @@ -135,7 +150,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "Tensor cu_seq_lens, Tensor cu_kv_comp, " "int batch_size, int page_size, " "int head_dim, int compress_ratio, " - "int next_n) -> ()"); + "int next_n, Tensor? new_tokens_per_seq=None) -> ()"); m.def( "compressor_prefill_reduction(" diff --git a/docs/dspark_bench_assets/p25fix_bench.py b/docs/dspark_bench_assets/p25fix_bench.py new file mode 100644 index 000000000000..d8bf9299bf34 --- /dev/null +++ b/docs/dspark_bench_assets/p25fix_bench.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +"""Round-style throughput bench: bs concurrent completions, JSON-line summary. + +Usage: p25fix_bench.py +Single-threaded asyncio client (no thread-pool: 1024 threads exhausts login +nodes and starves its own concurrency). Run node-local against the serve. +Mirrors the p25 rounds methodology: poetry (single prompt fan-out) and arena +(question cycle), bs in {512, 1024} global, 5 reps, max_tokens 768, temp 0.7; +metric is aggregate completion tokens / wall seconds. A leading warm rep per +(dataset, bs) is tagged rep=-1 so steady-state reps are comparable. +""" +import asyncio +import json +import os +import sys +import time + +import aiohttp + +HOST, PORT, TAG, SUMMARY = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4] +RUNS = "/lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs" +MODEL = (sys.argv[5] if len(sys.argv) > 5 else + "/lustre/fsw/coreai_comparch_trtllm/laliao/llm-models/DeepSeek-V4-Pro-DSpark") +URL = f"http://{HOST}:{PORT}/v1/completions" + +poetry = open(f"{RUNS}/poetry_prompt.txt").read().strip() +arena = [json.loads(l)["prompt"] for l in open(f"{RUNS}/arena_questions.jsonl")] + + +async def one(session, prompt): + body = {"model": MODEL, "prompt": prompt, + "max_tokens": 768, "temperature": 0.7} + try: + async with session.post(URL, json=body) as r: + d = await r.json() + return d["usage"]["completion_tokens"], len(d["choices"][0]["text"]) + except Exception: + return 0, 0 + + +async def bench(session, dataset, prompts, bs, rep): + t0 = time.time() + results = await asyncio.gather( + *[one(session, prompts[i % len(prompts)]) for i in range(bs)]) + wall = time.time() - t0 + toks = sum(r[0] for r in results) + chars = sum(r[1] for r in results) + fails = sum(1 for r in results if r[0] == 0) + line = json.dumps({ + "tag": f"{TAG}_{dataset}", "bs": bs, "rep": rep, + "completion_tokens": toks, "wall_s": round(wall, 2), + "out_tok_per_s": round(toks / wall, 1), + "chars_per_tok": round(chars / max(toks, 1), 2), "fails": fails, + }) + with open(SUMMARY, "a") as f: + f.write(line + "\n") + print(line, flush=True) + + +async def main(): + timeout = aiohttp.ClientTimeout(total=1800) + conn = aiohttp.TCPConnector(limit=1100) + async with aiohttp.ClientSession(timeout=timeout, connector=conn) as s: + bs_list = tuple( + int(x) for x in os.environ.get("BENCH_BS_LIST", "512,1024").split(",")) + for dataset, prompts in (("poetry", [poetry]), ("arena", arena)): + for bs in bs_list: + await bench(s, dataset, prompts, bs, -1) # warm rep + for rep in range(5): + await bench(s, dataset, prompts, bs, rep) + print("BENCH-DONE", TAG, flush=True) + + +asyncio.run(main()) diff --git a/docs/dspark_bench_assets/poetry_prompt.txt b/docs/dspark_bench_assets/poetry_prompt.txt new file mode 100644 index 000000000000..3fdbcd4a147b --- /dev/null +++ b/docs/dspark_bench_assets/poetry_prompt.txt @@ -0,0 +1 @@ +Write a long, original free-verse poem in English exploring memory, rivers, glass, and the passage of time. Use vivid, surprising imagery; vary the rhythm and stanza length; avoid repeating phrases or returning to the same image twice. Let the poem wander through at least eight distinct scenes. diff --git a/docs/dspark_bench_assets/postfix_pro_table.json b/docs/dspark_bench_assets/postfix_pro_table.json new file mode 100644 index 000000000000..df7fe7807dae --- /dev/null +++ b/docs/dspark_bench_assets/postfix_pro_table.json @@ -0,0 +1,302 @@ +{ + "token_counts": [ + 96, + 128, + 192, + 256, + 320, + 384, + 640, + 768 + ], + "step_time_ms": [ + 5.332674, + 7.125582, + 19.072141, + 22.564258, + 30.341472, + 32.343115, + 54.299678, + 61.859097 + ], + "fixed_overhead_ms": 39.296781, + "batch_sizes": [ + 32, + 64, + 128 + ], + "batch_overhead_ms": [ + 0.0, + 3.221512, + 9.122014 + ], + "_meta": { + "lookup": "interp", + "encoding": "decomposed", + "note": "step_time_ms is theta(M) only -- the trimmable term. The non-trimmable floor is fixed_overhead_ms + batch_overhead_ms[bs]. Do not read step_time_ms as whole-step time.", + "token_unit": "total_verify_tokens = num_requests * (verify_len + 1); the bonus token is included", + "cells": [ + { + "batch_size": 32, + "verify_len": 2, + "total_verify_tokens": 96, + "step_time_ms": 44.629455, + "num_samples": 114, + "p10_ms": 42.656446, + "p90_ms": 45.921993 + }, + { + "batch_size": 32, + "verify_len": 3, + "total_verify_tokens": 128, + "step_time_ms": 46.422362, + "num_samples": 252, + "p10_ms": 44.540524, + "p90_ms": 48.05913 + }, + { + "batch_size": 32, + "verify_len": 5, + "total_verify_tokens": 192, + "step_time_ms": 58.368921, + "num_samples": 71, + "p10_ms": 56.812048, + "p90_ms": 60.264587 + }, + { + "batch_size": 64, + "verify_len": 2, + "total_verify_tokens": 192, + "step_time_ms": 61.590433, + "num_samples": 87, + "p10_ms": 60.002756, + "p90_ms": 63.626337 + }, + { + "batch_size": 64, + "verify_len": 3, + "total_verify_tokens": 256, + "step_time_ms": 65.08255, + "num_samples": 77, + "p10_ms": 63.558149, + "p90_ms": 67.810392 + }, + { + "batch_size": 64, + "verify_len": 4, + "total_verify_tokens": 320, + "step_time_ms": 72.859764, + "num_samples": 54, + "p10_ms": 71.078515, + "p90_ms": 74.275923 + }, + { + "batch_size": 64, + "verify_len": 5, + "total_verify_tokens": 384, + "step_time_ms": 74.861407, + "num_samples": 70, + "p10_ms": 72.710967, + "p90_ms": 76.689982 + }, + { + "batch_size": 128, + "verify_len": 2, + "total_verify_tokens": 384, + "step_time_ms": 80.761909, + "num_samples": 81, + "p10_ms": 78.371763, + "p90_ms": 82.598686 + }, + { + "batch_size": 128, + "verify_len": 3, + "total_verify_tokens": 512, + "step_time_ms": 90.330839, + "num_samples": 47, + "p10_ms": 87.460041, + "p90_ms": 92.462635 + }, + { + "batch_size": 128, + "verify_len": 4, + "total_verify_tokens": 640, + "step_time_ms": 102.718472, + "num_samples": 46, + "p10_ms": 99.931121, + "p90_ms": 105.115652 + }, + { + "batch_size": 128, + "verify_len": 5, + "total_verify_tokens": 768, + "step_time_ms": 110.277891, + "num_samples": 53, + "p10_ms": 108.377361, + "p90_ms": 113.493729 + } + ], + "fit": { + "max_abs_residual_ms": 0.0, + "max_rel_residual": 0.0, + "additive_tolerance": 0.02, + "riser_tolerance": 0.01, + "min_riser_ms": 0.728598, + "typical_step_ms": 72.859764, + "warnings": [] + }, + "block_size": 5, + "min_verify_len": 1, + "modal_batch_size": 128, + "warmup_steps_discarded": 128, + "statistic": "median", + "source": "iter_log", + "logs": [ + "/lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs/protbl_744692_filtered.log", + "/lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs/protbl_744693_filtered.log" + ], + "recommended_confidence_verify_len_tiers": [ + 1, + 3, + 5 + ], + "recommended_tiers_note": "Nothing reads this from the table file. Copy it into DSparkDecodingConfig.confidence_verify_len_tiers, which is also what sizes the captured CUDA graph set.", + "diagnostics": { + "step_time_spread_by_batch_size": { + "32": 0.307856, + "64": 0.507969, + "128": 0.553581 + }, + "best_step_time_spread": 0.553581, + "profitability_probe": [ + { + "batch_size": 32, + "acceptance_rate": 0.3, + "chosen_verify_len": 3, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 32, + "acceptance_rate": 0.5, + "chosen_verify_len": 3, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 32, + "acceptance_rate": 0.7, + "chosen_verify_len": 3, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 32, + "acceptance_rate": 0.8, + "chosen_verify_len": 3, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 32, + "acceptance_rate": 0.9, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 32, + "acceptance_rate": 0.95, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 64, + "acceptance_rate": 0.3, + "chosen_verify_len": 1, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 64, + "acceptance_rate": 0.5, + "chosen_verify_len": 1, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 64, + "acceptance_rate": 0.7, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 64, + "acceptance_rate": 0.8, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 64, + "acceptance_rate": 0.9, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 64, + "acceptance_rate": 0.95, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 128, + "acceptance_rate": 0.3, + "chosen_verify_len": 1, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 128, + "acceptance_rate": 0.5, + "chosen_verify_len": 1, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 128, + "acceptance_rate": 0.7, + "chosen_verify_len": 3, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 128, + "acceptance_rate": 0.8, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 128, + "acceptance_rate": 0.9, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 128, + "acceptance_rate": 0.95, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + } + ], + "trims_somewhere": true + } + } +} diff --git a/docs/dspark_bench_assets/pro31_iterlog.yaml b/docs/dspark_bench_assets/pro31_iterlog.yaml new file mode 100644 index 000000000000..36ef5e7d8911 --- /dev/null +++ b/docs/dspark_bench_assets/pro31_iterlog.yaml @@ -0,0 +1,21 @@ +speculative_config: + decoding_type: DSpark + max_draft_len: 5 + speculative_model: /lustre/fsw/coreai_comparch_trtllm/laliao/llm-models/DeepSeek-V4-Pro-DSpark + enable_confidence_scheduling: false +moe_config: + backend: TRTLLM +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.5 +enable_attention_dp: true +cuda_graph_config: + enable_padding: true + batch_sizes: + - 1 + - 8 + - 32 + - 64 + - 128 +disable_overlap_scheduler: false +print_iter_log: true diff --git a/docs/dspark_bench_assets/pro_sched_arm.yaml b/docs/dspark_bench_assets/pro_sched_arm.yaml new file mode 100644 index 000000000000..370a68531eb2 --- /dev/null +++ b/docs/dspark_bench_assets/pro_sched_arm.yaml @@ -0,0 +1,28 @@ +speculative_config: + decoding_type: DSpark + max_draft_len: 5 + speculative_model: /lustre/fsw/coreai_comparch_trtllm/laliao/llm-models/DeepSeek-V4-Pro-DSpark + enable_confidence_scheduling: true + enable_ragged_verify: true + confidence_verify_len_tiers: + - 2 + - 3 + - 4 + - 5 + confidence_sps_table_path: /lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs/postfix_pro_table.json +moe_config: + backend: TRTLLM +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.5 +enable_attention_dp: true +cuda_graph_config: + enable_padding: true + batch_sizes: + - 1 + - 8 + - 32 + - 64 + - 128 +disable_overlap_scheduler: false +print_iter_log: true diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md new file mode 100644 index 000000000000..d7d0e739c1f8 --- /dev/null +++ b/docs/dspark_sts_program.md @@ -0,0 +1,1090 @@ +# DSpark STS 校准实验程序(任务 #30) + +> 活文档:计划、进度、结果、分析都记在这里,随实验推进更新。 +> 分支 `dspark-p25-device-window`;负责段:2026-08-07 起。 + +## ⚠️ REVIEWER 导读(2026-08-08 07:00 快照,供另一 agent 协查) + +**现状一句话**:device 端窗口选择(#25,env `TLLM_DSPARK_DEVICE_WINDOWS=1` 门控,默认关)会静默腐蚀生成;两个 bug 已定位一个、修复一个;第二个已围至"序幕 × w≠S × 稳态多请求工况",正在用 at-scale 装配差分收网。**host 路径(PR 主线)经受住了所有压力测试,无恙。** + +**已定案**: +1. **Bug #1(已修,勿回退)**:`_apply_device_window_prologue` 对 kv_lens 的差量应为 `2×(w−S)`(host 装配公式 `kv_lens = num_cached + seq_lens_kv` 两项都含本步窗口),原单倍差量使 indexer K-cache 写槽(slot_mapping_fp8)错位。经进程内 A/B 张量差分逐位验证;单测 `test_prologue_kv_pairing_matches_host_staging` 锁定公式。 +2. **旧实验结论修正**:此前 device "+63.9%" 吞吐是腐蚀假象(垃圾 token 易接受 + 零 EOS 满批不排水);host 动态裁剪四轮复现 ≈ 0±3%(校准+修表后)。 + +**未定案(协查重点)**: +- Bug #2:序幕激活(w≠S)+ 多并发 → 逐请求 census 63-64/64 全腐蚀,**与温度/采样图变体无关**(单探针曾误导出变体理论,已被 census 推翻)。prepare 级 21 张量 A/B 在 ramp 期小批步上全等;稳态步未采样(at-scale A/B 进行中)。 +- 已排除(每项有受控实验):重写机制(恒等/hybrid)、共享 ragged 消费方(host+steep 净)、sampler ride(仪表直证)、接受越窗(不变量仪表 0 违规)、V2 rewind(窗口不变代数)、图内采样计算(force-argmax)、逐 token 采样参数 H2D(skip 实验)、双趟捕获分离(交错捕获)、eager(该模式本就不支持序幕,host 标量按 S 直喂 kernel)。 +- **建议协查视角**:① 序幕在 padded_bs>n_real / 大 bucket / 批组成变化步的数学(`_apply_device_window_prologue` 全部下标与边界);② 序幕写入与图内消费方的契约面(dsa.py `apply_device_ragged_layout`/`on_update_kv_lens` 链);③ 跨步状态链(store→next staging 的张量时序)。 +- 完整实验叙事(20+ 受控判别)见下方进度日志;**census 是唯一可信文本判据**(单请求探针会被幸存请求误导)。 + +**本分支的调试支架(全部 env 门控、默认关闭、定案后剥离,勿 review 为正式代码)**: +`TLLM_DSPARK_DEVICE_IDENTITY`(1/steep,序幕恒等/极端窗口)、`TLLM_DSPARK_HOST_STEEP`(host 发布非均匀窗口)、`TLLM_DSPARK_FORCE_EAGER`+`EAGER_OK`、`TLLM_DSPARK_AB_DUMP`+`AB_AFTER`(装配 A/B 差分)、`TLLM_DSPARK_RIDE_LOG`(ride/accept-bound 仪表)、`TLLM_DSPARK_FORCE_ARGMAX`、`TLLM_DSPARK_SKIP_SAMPLING_H2D`、`TLLM_DSPARK_SKIP_GREEDY_CAPTURE`/`SKIP_ADVANCED_CAPTURE`/`INTERLEAVED_CAPTURE`(捕获趟操控)。涉及文件:model_engine.py、spec_sampler_base.py、interface.py。 +**复现最小配方**:DEP8 serve + `p25_round_arm.yaml`(sts_v1+sps_frac_v2b+tiers[2,3,4,5])+ `DEVICE_WINDOWS=1 DEVICE_IDENTITY=steep` + 64 并发固定 prompt(`dspark-runs/percreq_payload.json`)→ census 全腐蚀;去掉 DEVICE_WINDOWS 或改 HOST_STEEP → 全净。 + +## 总目标 + +1. 采集并拟合 STS 校准(此前所有实验都是裸 sigmoid,未校准); +2. 在校准 + 修复表的条件下,验证细刀(tiers [2,3,4,5])下 host/device 对 notrim 的收益在 **bs 64/128 per-rank(agg DEP8)** 是否**跨节点多轮稳定复现**; +3. 若复现:**搞清楚 device 为什么高于 host**; +4. draft-len 维度:draft=5 的三模式对比,以及 "draft5+ragged vs 最优固定短 draft(3/4/5)+满块" 的零假设检验。draft 6 已砍(checkpoint 无第 6 块权重,用户确认不看)。 + +## 背景:为什么要做这一轮(初测结果及其瑕疵) + +2026-08-07 凌晨,细刀三臂初测(poetry temp0.7 out512,v2 表): + +| bs(per-rank) | notrim | host | device | +|---|---|---|---| +| 512(64) | 11269 | 11319 (+0.4%) | 14537 (+29.0%) | +| 1024(128) | 11300 | 12236 (+8.3%) | **18524 (+63.9%)** | + +调度统计:notrim trim 0% accept 1.162;host trim 33.7% accept 1.077;device trim 49.7% accept_stat 2.032(可疑)。 + +三个瑕疵(本轮逐一修正): +1. 无 STS 校准(裸 sigmoid); +2. v2 表拟合被 `compress_to_risers` 截断到 8 节点:θ(512) 实测 33.9ms 被插值低估 ~9%,384→480 的真实陡坡被抹平; +3. device 段与 notrim/host 段跑在不同节点(晚绑定切换),跨节点混杂未消。 + +同期对照(粗刀 [1,3,5],同负载):host −8.6%~+3.8%、device ≈0 —— **tier 粒度是收益解锁器**(#19 论点的实验证据)。 + +## 实验步骤 + +### Step 1:STS 采集与拟合 【✅ 完成 10:18】 +- **结果**:`sts_v1.json`,温度 = [2.512, 1.585, 1.122, 1.995, 1.413];ECE 前 [0.130, 0.145, 0.115, 0.083, 0.061] → 后 [0.074, 0.081, 0.064, 0.047, 0.030],**5/5 位置改善,约减半**; +- **解读**:全部温度 >1 —— 裸 sigmoid 在所有 draft 位置**过度自信**(位置 0 最重,T=2.51)。意味着此前未校准实验的存活概率被系统性高估 → 裁剪偏保守的方向性偏差;位置间非单调(位置 3 T≈2.0)属 per-position 自由拟合的正常形态。 +- static 模式满块 serve(采集守卫强制无审查标签:无成本表、无裁剪),DEP8; +- 负载:poetry + arena 混合,temp 0.7,bs 512 global,各 2 遍; +- 拟合:`dspark_fit_sts.py`(与 SGLang 同方法:41 点对数网格 / 左到右冻结 / cumprod-ECE / 15 bins)→ `sts_v1.json`; +- 与 SGLang 的差异仅在配对管道(overlap 边界 → draft_seq_ring 双轴配对);拟合结果可互换。 +- 注:SGLang blog 实验的 `speculative_dspark_confidence_sts_path=None` —— 对照系统同样未校准,我们此前的未校准实验与 blog 同口径。 + +### Step 2:成本表修复 【✅ 完成】 +- 同样本重拟合,`--max-breakpoints 16` → `sps_frac_v2b.json`(16 节点); +- 不需要重新采集(截断发生在拟合层,原始样本完整)。 + +### Step 3:复现性轮次(核心)【等 STS 就绪自动开跑】 +- **4 节点并行,每节点一轮独立实验**(不做同节点编排,跨节点复现本身就是判据——用户指示); +- **每轮 3 次冷启动 serve**(notrim → host → device)——每臂重启,杜绝 reuse/分配器/autotuner 遗留污染(用户指示); +- 每臂:**poetry 与 arena 各做一次完整网格**(用户指示):每数据集 bs {512, 1024} global(64/128 per-rank)× 5 遍 × output 768 —— 每臂测量 ~17 分钟,每轮 ~2 小时; +- 配置:`sts_v1.json` + `sps_frac_v2b.json` + tiers [2,3,4,5],DEP8,max_batch_size 128; +- 三臂环境变量:notrim=`TLLM_DSPARK_FORCE_BUDGET_FRAC=1.0`;host=无;device=`TLLM_DSPARK_DEVICE_WINDOWS=1`; +- **判决**:4 轮的 (host−notrim)/notrim 与 (device−notrim)/notrim,方向一致且轮间散布 ≪ 效应 → 复现成立。 + +### Step 4:根因分析(device 为何 > host)【门禁:Step 3 复现】 +- **H3**(优先):accept 统计审计——device accept_stat 2.03 ≈ 其平均窗口("每请求接满窗"特征)。三种解释与判别器: + - (a) **真机制**:device 按当步置信度排序,留下的正是会被接受的位置 → accept≈窗口本就是好调度器的形态(设计意图); + - (b) **统计口径缺陷**:ride 记账(token 窗 −1)与 drafted 口径错位 → 判别:算术闭环(门禁 3)——client 侧 tokens/steps 独立推出 1+accept,与调度器 ride 值对表;吞吐收益若与 steps 降幅成比例,则统计为真; + - (c) **接受判定损坏**(接了不该接的)→ 判别:**temp-0 逐 token 等价**——同 prompt 下 device 臂 vs static 满块臂输出必须逐字一致(投机解码无损性);不一致即判定损坏。GSM8K 已过但 bs 小,需在 bs128 复核; + - 判别顺序:先 (b)(轮次数据自带)→ 再 (c)(等价试验)→ 都排除则接受 (a),交由 H2 分离"预算更狠"与"排序更新鲜"。 +- **H2**(判决性):裁深匹配差分——host + FORCE_BUDGET_FRAC≈device 实际裁深(~0.5)vs device,拆开"预算更狠"与"排序更新鲜"; +- **H4**:槽位新鲜度——长输出(低流转)下差距是否收窄。 + +### Step 5:draft-len 维度 【门禁:Step 3 复现】 +- 5a:draft=5,三模式对比:ragged-compact / cap-accept / static 满块,@ bs 64/128; +- 5b:零假设检验——**ragged@draft5 vs 最优 static@draft{3,4,5}**:动态裁剪 ≈ 自适应短草稿,若 static-3 就能拿走大部分收益,调度机器的增量价值存疑; + - 前置:放宽 worker 的 `block_size == max_draft_len` 断言(截断用前 k 块)+ GSM8K@draft3 正确性验证; + - draft 6:已砍。 + +## 挂起项(本程序期间不动) +- #29 disagg 512/rank 解锁(gen 侧 V2+NIXL CUDA IMA 未解); +- PR 推送(6 commit 就绪,等用户侧 git 认证通道); +- #11/#19/#20 等。 + +--- + +# 进度日志 + +- **08-07 09:04** STS 采集首发:配置被拒(`enable_confidence_scheduling` 要求 `enable_ragged_verify=true`)。 +- **08-07 09:14** 修复(yaml ragged=true + static 环境变量)重发,735105 采集中。 +- **08-07 09:20** v2b 表重拟合完成:16 节点;θ(512)=33.9ms(v2 插值低估 9%);384→480 实测陡坡恢复(23.9→31.7ms)。 +- **08-07 09:25** 一次未经确认的抢跑被用户叫停,轮次与节点已回收;计划经 review 修订(每臂重启、reps=5/out=768、加 arena 对照、砍 draft6)后获准。 +- **08-07 09:30** 4 轮次节点排队(735828–31),派发器就绪;轮脚本带 STS+v2b 门禁,采集完成后自动开跑。 +- **08-07 09:40** 用户升级 arena 为完整实验(原为单格对照);轮脚本更新为 poetry+arena 双全网格,重派。 +- **08-07 09:16→09:42** STS 采集第 2 次发射死于配置连环套(24 分钟后才发现——盯梢失职,已换 120s 紧密监视 + serve 日志错误早检):`enable_confidence_scheduling` 要求 ragged=true,ragged=true 要求 cost table 或 sts path,而采集守卫拒绝 cost table。**解法:恒等 STS 垫门**(`sts_identity.json`,全 1.0 温度 = 裸 sigmoid,零偏置)满足校验;第 3 次发射 09:42 起。节点门禁同步放宽(≥1h,采集只需 ~45 分钟)。 +- **08-07 09:46** 第 3 次发射死于 **GPU 显存不足**:735105 是此前 disagg gen IMA 事故节点,`kill -9` 后 CUDA 上下文僵死(D-state),274GB/卡 无法回收——**IMA 事故节点即废节点**,只能整作业回收。教训:晚绑定节点选择两度肇事(交叉误杀 + 复用废节点),STS 战役改为**显式钉死节点**。735105 已 scancel。 +- **08-07 09:51** 第 4 次发射:钉死 735104(disagg ctx 节点,验证 0 进程 0 显存,3:58 余量),120s 紧监视 + serve 错误早检在岗。 +- **08-07 10:18** STS 采集+拟合完成(serve 就绪 10:16,4 段混合负载 ~2.5 分钟,168 rank 文件,拟合 exit=0);`sts_v1.json` 落盘,轮次门禁 10:18:58 放行,round1(735830)/round2(735831) notrim 臂开跑。round3/4 等 735828/829 出队后自动派发。 + +## 数据合理性门禁(每轮数据落地即检查;异常 → 当场 debug → 记录在下方 debug 日志) + +1. **调度自洽**:notrim 臂 trim 必须 ≈0;各臂 delivered/ceiling 与 trim_ratio 一致;fallback/graph_eager 计数不异常膨胀; +2. **接受合理性**:accept_len ≤ 平均窗口;**accept ≈ 满窗是饱和红旗**(现存疑点:初测 device accept_stat 2.03 ≈ 其平均窗,即 H3);accept 与 trim 的方向关系合理(裁得多不应接得更多); +3. **算术闭环**:吞吐 ↔ steps_total ↔ (1+accept)×并发/步时 三者量级互恰;completion_tokens 接近 bs×max_tokens(poetry EOS 折损 ~7% 内); +4. **统计卫生**:rep 间离群(>2σ)标记并查因;首 rep 爬坡污染;轮间方向翻转即触发 debug 而非平均掉; +5. **臂间守恒**:三臂 completion 总量应近似(同负载同上限),显著偏差 = 生成行为被改变的信号(优先查 rewind/接受判定,参考 08-07 抓到的 cap-accept capture bug 的教训)。 + +任何一条触发:暂停下结论 → 定位(优先顺序:统计口径 bug > 测量方法 bug > 真实机制)→ 结论与证据写入下方 debug 日志,修复后视影响面决定重跑范围。 + +# 实验结果 + +## Step 3 轮次数据(吞吐 tok/s,mean 去 rep0 爬坡,n=5) + +### notrim 基线(四轮全过门禁 1:trim 全 0.0;accept 1.78/1.73/1.77/1.74 轮间高度一致) + +| 轮(节点) | poetry bs512 | poetry bs1024 | arena bs512 | arena bs1024 | +|---|---|---|---|---| +| round1 (735830) | 11505 | 12520 | 16409 | 13761 | +| round2 (735831) | 11430 | 12128 | 15440 | 13751 | +| round3 (735828) | 11591 | 12136 | 15142 | 13626 | +| round4 (735829) | 11198 | 12167 | 16246 | 13755 | +| **轮间散布** | ±1.7% | ±1.6% | ±4.2% | ±0.5% | + +注:arena bs512 测量窗最短(~24s),散布最大;作对比分母时以同轮同臂配对为准,不跨轮混用。 + +### host 臂(完整,四轮全过门禁:trim 22.4-22.7%、accept 1.66-1.67 逐位一致) + +| 轮 | poetry bs512 | poetry bs1024 | arena bs512 | arena bs1024 | +|---|---|---|---|---| +| round1 | +3.2% | −0.1% | +1.0% | −2.4% | +| round2 | +0.7% | +1.9% | −1.0% | −0.1% | +| round3 | −0.7% | +1.3% | +1.6% | −1.4% | +| round4 | +5.7% | +2.2% | −3.7% | −5.6% | +| **均值** | **+2.2%** | **+1.3%** | **−0.5%** | **−2.4%** | + +### device 臂(数据作废——生成损坏,见 Debug 日志;吞吐 +50%~+55% 为假象) + +device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8%、accept ≈2.06 ≈ 平均窗口(饱和)。 + +# Debug 日志 + +## 【08-07 12:30】device 臂生成损坏实锤(门禁 5 触发 → 活体探针定罪) + +**触发**:round1 device poetry bs512 的 completion_tokens 连续多 rep **恰好 = 393,216 = 512×768(100.0% 上限,零 EOS)**,而同轮 notrim/host 均 ~96%(自然 EOS 折损)。无损投机解码不应改变文本分布——只有 device 臂 EOS 行为漂移。 + +**探针 1**(temp-0 事实问答,round1 device serve):`' Paris"\n '`,正常终止——排除全面乱码。 + +**探针 2(定罪)**(真 poetry prompt,temp 0.7,max_tokens 768,与批负载并行):`finish_reason: length`,**768 token 生成后全文只有 2 字符 `S 读越界/垃圾 KV,w ride 窗口` = 图内接受越过本行布局窗口的现行证据;`ride 窗口=0` = ride 缺失导致 rewind 回退 S 快照的现行证据。738957 已带此仪表重启,steep+graph,收割中。 + +## 【08-07 23:55】accept-bound 收割:零越窗违规;单请求净;转 bs 扫描找最小复现 + +- **零次"接受越窗"**:整段脏负载中,图内接受从未超过本行布局窗口——接受计数与布局一致,排除"越窗读写"类; +- **ride 零值波全部出现在混合/过渡步**,那里回退 S 恰好等于该步布局(无序幕),行为正确; +- **单请求(batch=1)探针完美干净**(987 字符连贯诗歌 temp0)——但注意 batch=1 时 steep 转移无操作(w≡S),判别力≈恒等; +- **推论收窄**:接受计数对、装配张量对、越窗不存在 → 腐蚀在"接受的 token 内容"或"跨请求 KV/状态污染",且需要多请求才触发; +- **过夜实验(738957)**:bs 扫描 {16,32,64,128,256,512} × temp0 × 每档文本探针——找腐蚀出现的最小批量,作为明天内核级挖掘的最小复现。 + +## 【08-08 00:55】当夜最终收口:残余 bug 锁定"advanced-sampling 图变体 × w≠S" + +- **bs 扫描(temp0)判决:2×修复后 greedy 路径全净**——bs 16/32/64/256/512 探针全部连贯(~4 chars/tok;128 的 2.4 是探针踩中 ramp 的噪声)。greedy 完整走 verify+draft+commit 管线 → **布局/KV/接受/提交机器全部无罪(修复生效)**。 +- **同一台 serve、同一份代码:temp0.7 负载 99.3% 撞顶(脏)**——残余腐蚀是**采样路径特异**的。temp0 replay greedy 图变体;temp0.7 replay **advanced-sampling 图变体**(独立捕获趟)——嫌疑收窄到该变体的捕获/绑定与 w≠S 的交互。 +- **ride 陈旧被证明良性**:ride=0+陈旧 py_verify_len 的波在两种温度下都大量出现(温度不相关),且 V2 管理器的两项 rewind(`py_rewind + max(reserved−runtime,0)`)总和 = `5−acc`,**与窗口无关**——rewind 正确性不依赖 ride。ride 只喂统计。 +- **明晨第一刀(廉价决定性)**:temp0.7 + w≠S + **禁用 advanced-sampling 捕获**(非 greedy 批将按注释回退 eager)。净 → bug 在 advanced 变体的捕获路径(对比两趟捕获的差异:`_force_non_greedy_for_capture`→populate→`_sample_tokens_for_batch` 非贪婪分支的捕获期绑定,重点:`spec_metadata.temperatures` 的对象身份/持久性、flashinfer `sampling_batch_spec_dec_one_model` 的 torch.compile 捕获);脏 → eager×0.7×w≠S 复现,范围另定。 +- 当夜战果:bug#1(kv_lens 2×)修复+验证;greedy 全尺寸验证干净;残余范围从"整个 device 路径"收窄到"一个图变体的捕获交互"。 + +## 【08-08 01:25】Bug #3 干净定罪:advanced-sampling 图变体 × device 窗口 + +- **禁 advanced 捕获 + temp0.7 = 干净**——但混杂(非贪婪回退 eager → 序幕不跑,等价 host 路径),仅作旁证; +- **去混杂终验(738957)**:`temperature=0.0001`(分布≈argmax;temp0 同负载 = 完美诗歌)→ **2 字符垃圾**。同 serve 同分布,唯一变量 = 路由到 advanced-sampling 图变体。**定罪成立**:greedy ragged 图 + 序幕 = 净;advanced ragged 图 + 序幕 = 毁;advanced ragged 图 + host 装配(host+steep@0.7)= 净。 +- **三元交互**:腐蚀 = advanced 图变体 × 序幕改写(w≠S)× 多请求。变体间的捕获差异只有 `_sample_tokens_for_batch` 非贪婪分支(flashinfer `sampling_batch_spec_dec_one_model`,torch.compiled)+ is_all_greedy 条件代码——明日在此范围内定位(候选:per-token 采样参数缓冲的捕获期对象绑定/装配路径、compiled kernel 的捕获期特化、seed/offset 语义;注意 host 装配可走通 → 嫌疑聚焦在"该分支读取了某个序幕没有重写、且按 S 装配的逐 token 缓冲——温度同质时值相等的推理可能有漏洞,如 dummy/probe 行的哨兵值在 w 布局下错位到真实行")。 +- **明日计划**:① 在 advanced 分支加 replay 后 D2H dump(sampled vs argmax @1e-4;temps/top_ks 实际值),一步定位;② 修复;③ graph steep 0.7 探针转净 → 真 device 吞吐网格 + GSM8K 大并发 → 剥支架、整理提交。 + +## 【08-08 01:30】变体内二分:捕获的算子无罪 → 嫌疑收缩到双趟捕获机制本身 + +- **force-argmax(捕获前置生效)+ advanced 变体 + 0.7 = 仍脏**(0.17 chars/tok)。此时 advanced 图的算子内容与 greedy 图完全一致(target argmax 已强制;draft 本就 argmax——`wants_advanced_draft_sampling` 需要 rejection,而 `use_rejection_sampling=False`);逐 token 采样参数缓冲同质中和。 +- **推论**:greedy 图净 / 计算等价的 advanced 图脏 ⇒ 差异在**双趟捕获机制**:图池共享/别名、第二趟捕获对共享持久状态的副作用、runner 对两变体的 key/元数据簿记——与序幕(w≠S)交互的面。 +- **明晨作战序列**:① 读 runner 的双变体捕获簿记(graph pool、shared_static_tensors、graph_metadata/outputs per key)找 w≠S 敏感的共享面;② 对照实验:只捕获 advanced 趟(跳过 greedy 趟)+ temp0.7 → 若净,则是"两趟共存"的交互;仍脏则是 advanced 趟自身的捕获环境(force_non_greedy 时的 populate/warmup 状态);③ 定位后修复 → 验证三件套 → 吞吐 + GSM8K 大并发 → 剥支架。 +- 注:所有调试开关(IDENTITY/FORCE_EAGER/EAGER_OK/AB_DUMP/RIDE_LOG/HOST_STEEP/SKIP_ADVANCED_CAPTURE/FORCE_ARGMAX/SKIP_GREEDY_CAPTURE)+ 仪表代码在修复定案后一并剥离;`kv_lens 2×` 修复与单测保留。 + +## 【08-08 02:00】当夜终判:双趟捕获共存被点名 + +- **只捕 advanced 趟(跳过 greedy 趟)+ steep + 0.7 = 大幅转好**(2.0 chars/tok,连贯语料延续;双趟时 0.01-0.17 垃圾)——**腐蚀 = 双趟捕获共存 × 序幕(w≠S)**。单独任一趟的图都(基本)健康;两趟共存时 temp0.7(pass-2 图)被毁。 +- 机制候选(明晨深挖):共享图池/attention 工作区在两趟重复捕获同形状时的指针/别名交互(代码内 dynamic-draft-len 的已知同类患:"resize 使先捕图的指针失效");pass-1 捕获对持久状态的遗留改变 pass-2 捕获期的 Python 条件求值;worker 持久状态(_ctx_len/_kv_windows/confidence)跨趟的捕获期录制差异。 +- 明晨序列:① adv-only 严格验证(完整网格 + completion 指纹,2.0 的余量需排除内容密度混杂);② runner 双趟簿记深读(shared pool、workspace resize、graph_metadata);③ 修复(候选:ragged bucket 仅在所需变体捕获 / 趟间池隔离 / max-shape 前置捕获同款 workaround);④ 验证三件套 → 吞吐 + GSM8K 大并发 → 剥支架。 +- morning 节点已排队:739289/739290(4h × 2)。 + +## 【08-08 02:30-04:45】参照指纹 + 逐请求普查:adv-only "转好" 是抽样错觉 + +- **干净参照线确立(host+steep@0.7 同网格)**:completion **96.6-96.7%**、accept 1.06、吞吐 9.9-12.1K——与 notrim 的 96% 一致。 +- **adv-only device 对照**:completion 99.0-99.8%、accept 2.14、吞吐 12.8-16.2K——**仍系统性偏移,腐蚀未消**。 +- **逐请求普查仪表(64 并发独立请求、逐个存文本)**:adv-only device @0.7 = **64/64 全腐蚀**(chars/tok 中位 ~0.03)。"单趟转好"“探针 3.67”均为抽样错觉(探针骑行时活跃请求少)。**双趟共存理论降级为噪声;回归主线:advanced 图变体(temp>0)× 序幕 = 全毁,greedy 图 = 净。** +- force-argmax 普查(判"eager 非贪婪主机侧代码 vs 图内 advanced 采样计算")因节点到期未完成;739900/739901 已排队,落地重跑。 +- **普查仪表沉淀**:`percreq_payload.json` + 64 并发 curl + chars/tok 分类——今后所有验证以此为准(单探针不再作为判据)。 + +## 【08-08 06:45-07:45】重装配差分在规模下判死;转只读不变量审计 + +- **at-scale A/B(AB_AFTER=150)**:跨过阈值后开火即崩——executor 死于与猎手 v1 相同的响应管道错误(`'int' object is not iterable`)。**结论:重装配式差分只能在安静的 ramp 步上活着,规模下对请求态的突变不可控,该路线放弃**(其 ramp 步"21 张量全等"的判决仍有效,但只覆盖小批步)。 +- **新仪表(只读不变量审计,`TLLM_DSPARK_INVARIANT_LOG=1`)**:序幕末尾 D2H 少量派生量,对照 host-with-w 的解析公式验证不变量(首个:`kv_lens[r] == (max_beam_num_tokens−1) + 2×w_r`),零突变、零集合通信、零重装配——任意规模可常开。违规打印现场 (row, past, w, S, n_real, padded_bs, bucket);每 100 步打检查点(正向确认仪表存活)。 +- 协查提示:kv 不变量若稳态全绿,后续只读不变量按序追加:anchor 对位(`input_ids[qo[r]] == next_new_tokens[slot,0]`)、draft 对位、position 基。序幕输出面逐项排掉后,嫌疑收缩到图内消费方与跨步状态链。 + +## 【08-08 07:35-08:10】kv 不变量稳态全绿;扩展不变量 × 双线并行(steep 猎场 + identity 对照) + +- **kv 不变量:200+ 稳态步 0 违规**,检查点机制同时正向确认仪表存活——序幕 kv 输出面(修复后)在规模下逐请求正确。 +- **第二批只读不变量上线**(steep,739901):`kvoff == new_lens − w`、anchor 对位(`input_ids[qo[r]] == next_new_tokens[0,slot]`)、draft0 对位、`prev_pos_idx/off` 首尾 token 对位——覆盖序幕的全部 token 级输出。 +- **仪表对照组**(identity,739900):同一套不变量在 w=S 下应全绿——排除仪表自身假阳性;若对照组报违规则仪表公式有错。 +- 若 steep 全绿:第三批备选不变量 = row maps 解析重算(req_idx/kv_correction 是 w 的纯函数)、`seq_lens_cuda == w`、`gen_token_repeats == w`、pad 行窗口 == pad_len;全绿穿透后嫌疑正式移交图内消费方与跨步链。 + +## 【08-08 08:15-08:40】第二批全绿 + 对照全绿;batch-3 在跑;跨行分歧检测器上线 + +- **steep 第二批不变量:300+ 步 0 违规**(kvoff/anchor/draft0/prev_pos 首尾对位)——序幕 token 级输出面全部无罪;**identity 对照同样全绿**——仪表无假阳性,判定有效。 +- **batch-3(740252 加载中)**:row maps 解析重算、seq_lens==w、repeats==w、position 打包(行内位差)。 +- **新杀器(`TLLM_DSPARK_ROWDIV_LOG=1`,下个 serve 周期生效)**:census 的 64 个请求完全相同 → temp0 下所有行每步接受的 token 必须一致(窗口只改变验证量、不改变内容)。sampler 内纯 host 检查:行间 token 元组分歧即腐蚀现行,直接给出**首个分歧步 + 分歧行的窗口类别(w=2 vs w=6)**——回答"哪类窗口先烂",并把腐蚀定位从"输出面"推进到"图内效果"层。 + +## 【08-08 09:15-09:25】rowdiv serve 重发 ×2;守恒论证解释单请求永净;图内消费方静态围剿启动 + +- **rowdiv serve(740252)warmup 挂死**:日志冻结于 09:15:37,进程活但 20+ 分钟无进展 → 杀掉,双节点并行重发(740252→p25rowdiv2_serve.log,740251→p25rowdiv3_serve.log,同配置 steep+ROWDIV_LOG),first-ready-wins 收割流水线挂上(先 ready 者压 64 并发 temp0 census → grep `[rowdiv]`)。若二次 warmup 挂死 → 实锤 rowdiv 仪表自身与 warmup 步的交互问题。 +- **结构性论证(新)**:窗口守恒 sum(w)=sum(S) 意味着 **n_real=1 时数学上强制 w=S** —— 单请求下序幕是恒等变换。这解释了单请求探针为什么永净:不是运气,是结构。任何 w≠S 敏感的凶手**只可能**在多请求下现形,与全部观测(多并发 63-64/64 腐蚀、单请求净)精确吻合。 +- **静态围剿(并行 workflow)**:4 个读者分扫 ① host staging 面(_prepare_tp_inputs + spec_metadata.prepare)② attention+indexer 面(apply_device_ragged_layout 覆盖清单 vs MLA/索引器 replay 时全部读取)③ 图内 worker 面(fused verify+draft 的每个 gather/scatter 索引、confidence 装填索引、捕获期烘焙常量)④ 跨步链(host 簿记以 S 推进而 device 以 w 运行之处)。汇总产出"w 敏感 ∧ replay 消费 ∧ 序幕未重写"嫌疑排名 + 每个的最廉价判决实验。已知热点候选:`spec_metadata.gather_ids`(interface.py:478,drafting_loops 以它选 logits 行——若 DSpark 路径同类索引按 S 装填且未被序幕重写,则 w≠S 时采样/验证读错行,跨请求串染,恰好匹配全部现象)。 + +## 【08-08 10:15-10:50】BUG #2 根因定案:compressor 的逐请求新 token 数被捕获的 H2D 钉死在 S + +- **双线判决同时落地,互相咬合。** + - **rowdiv(动态,steep+temp0+64 相同请求)**:step#3-4 起行间 token 元组即分歧,同窗口类 (w=6) 行内容也互不相同,大批行迅速塌缩为 `(0,0,0)` 退化 token——腐蚀即时发生、按请求永久化,是"逐请求持久状态被写坏"的指纹。 + - **静态围剿(4+1 workflow,60 个缓冲逐一分类)**:唯一同时满足全部观测(多请求 only / identity 净 / temp0+force-argmax 也腐蚀 / 与捕获结构无关 / 序幕输出面逐项正确)的嫌疑就是 **`gen_new_tokens_per_seq_cuda`**。 +- **机制(全链路代码实证)**:host staging 在 `_sync_gen_tokens_per_seq`(deepseek_v4.py)把 host 形状split S 写入 pinned 缓冲并 H2D 到 `gen_new_tokens_per_seq_cuda`;该调用同时在捕获区 `on_update_kv_lens` 内执行,于是 **H2D memcpy 成为图内节点,每次 replay 从 pinned 缓冲重新加载当步的 S**——序幕改了也会被 replay 覆盖(序幕本来也没改它:`apply_device_ragged_layout` 写的是双胞胎 `seq_lens_cuda`/`gen_token_repeats_cuda`)。compressor kernel 算 `nn = new_tokens_per_seq[b]`、`sp = kv_len − nn`,而 `kv_len` 已被序幕修到 w 口径 → **`sp = past + (w−S)`**:Phase 1 把**持久分页 KV/score 递归状态**写错位(w>S 尾巴永远缺失;whost 不复现,原 +64% 判定为腐蚀假象(定案) + +修复后干净生成下的三臂对照(rounds 方法论,节点本地 asyncio 客户端,median tok/s,5 稳态 reps,warm rep 剔除;device=bia0047 / host=bia0048 / notrim=bia0123): + +| cell | notrim | host | device | host/notrim | device/host | +|---|---|---|---|---|---| +| poetry 512 | 12340 | 11840 | 11058 | −4.1% | −6.6% | +| poetry 1024 | 12912 | 12703 | 12385 | −1.6% | −2.5% | +| arena 512 | 16766 | 16221 | 16275 | −3.3% | +0.3% | +| arena 1024 | 13526 | 13636 | 13350 | +0.8% | −2.1% | + +- **主问题答案(#30 "若 device>host 复现,搞清楚为什么")**:干净生成下**不复现**——device 相对 host 在 −6.6%~+0.3% 区间(本轮未做节点对调,±3-5% 内视为噪声;poetry 512 的 −6.6% 或含真实序幕开销成分)。原 +64% 完全由腐蚀假象制造(重复特殊 token + 零 EOS + 满批不排水)。 +- host trim vs notrim 仍 ≈ 0±4%,与此前四轮判决(0±3%)一致。**B300 agg DEP8 上,该校准配置下 host/device 动态裁剪均无可复现吞吐收益**,与 #28(SGLang 原生动态在 B300 高 bs 亏损)闭环互证。 +- 量具备注:首版登录节点线程池客户端因线程耗尽产生 6× 失真(1998 tok/s),已废弃;节点本地 asyncio 客户端复现旧轮次量程(warm rep 11954 vs 旧轮 10.4-11.9k)。散在失败请求 <2%/格、三臂均摊,不影响中位数。 +- 至此 #30 全链闭环:STS 校准 → 稳定性轮次 → device>host 根因(= bug #1 kv_lens 双份 + bug #2 compressor 捕获 H2D)→ 修复 → census/GSM8K/吞吐三重验证。 + +### 【08-08 15:45】"为什么没有收益"定量分解(三臂日志计数器差分 + 成本表 + SGLang 原生对照) + +**方法**:三臂 serve 日志的 periodic 计数器(全部为累积量)按 bench 阶段做差分,得到每格的逐阶段 accept_len / trim / bucket_hist(每步实际重放的捕获桶);与 v2b 成本表的分解语义(θ(M) 可裁项 + 29.5ms 固定 + α(bs) 批开销)合成。notrim 臂无统计(DSparkRaggedStats 只在 confidence-scheduling 分支创建——仪表盲区,已记录),其 accept 由吞吐比值模型反解并与四轮存档交叉验证。 + +**每阶段实测(rank-0 切片,稳态区)**: + +| 阶段 | 臂 | accept_len | 落地 trim | 请求 trim | 主导桶 | +|---|---|---|---|---|---| +| poetry512 | host | 1.225 | 31.8% | 34.4% | 256@bs64 (61%) | +| poetry512 | device | 1.261 | 32.1% | — | 256@bs64 (57%) | +| arena512 | host | 2.534 | **5.8%** | 13.5% | 384@bs64 (=满块) | +| arena512 | device | 2.585 | 6.3% | — | 384@bs64 | +| poetry1024 | host | 1.208 | 33.1% | 35.4% | 256@bs64 (53%) | +| arena1024 | host | 2.566 | 7.6% | 16.3% | 384/432/320 混 | + +**四个成因,量化到每一格:** + +1. **可裁蛋糕本来就小(B300 成本结构)**:bs64/rank 满块步 ≈ 60.3ms,其中固定 29.5 + α(64) 6.9 = **60% 不可裁**;bs128 也有 47%。θ(M) 还有宽平台(θ(320)=θ(384)=23.85、θ(256)=θ(288)):**tier 5→4 的第一档裁剪 kernel 时间省 0ms**。首个有实际节省的档要砍到 256 桶(省 5.0ms = 8.3%/步)。 +2. **accept 损失 > 步时节省(poetry,主项)**:host 裁 32% token 拿到 8.3% 步时节省,但 accept 1.53(反解的 notrim 值,poetry1024 反解 1.45 交叉一致;四轮混合存档 1.73-1.78 与 arena≈2.55 的混合构成也吻合)→ 1.23,tokens/step −12%。净:0.88 × 1.09 ≈ 0.96 = **实测的 −4.1%**。模型在 poetry1024 同样闭合(−1.6%)。 +3. **机制税(arena,唯一贡献项)**:arena 上校准后的 planner 正确地几乎不裁(请求 13.5-16.3%,**经桶取整落地只剩 5.8-7.6%**——fill_bucket 把取整余量还给真实验证),但仍付信心头 + planner 集合 + ragged fit 的常数税 → arena512 −3.3%(arena1024 +0.8% 噪声内;该阶段 host 臂另有 222 步 eager 逃逸)。 +4. **device 额外的序幕开销**:device 的 lag-0 重排**确实提升 accept**(1.225→1.261,窗口直方图向 3 集中——机制生效),但序幕的十几个小 kernel + D2H ride 的每步开销吃掉了它,净 −2.5~−6.6% vs host。 + +**独立复核(#28,SGLang 自家实现原生跑 B300)**:bs=1 时动态 **+17.8%**(iter 16.7→13.2ms),bs 8-64 持平,bs128 **−21±4%**(accept 4.65→4.08 而 iter 反升 28.1→33.2ms),bs256 −4~−8%。收益只存在于 θ 陡峭的低占用 regime;高并发下 B300 的 θ 平台宽到裁剪换不回 accept——**与实现无关,是硬件成本结构**。SGLang blog 的 +10-20% 在 H200/H100 上 θ 份额高、平台窄,不可迁移。 + +**次要成因(可改进,但理论上限也有限)**:① 表的已知偏差(L-sweep 连 draft pass 一起缩,θ 边际被高估 → 过度裁剪);② serve 用 tiers [2,3,4,5] 而表自荐 [1,2,5](密集档全躺在平台上,缺低 accept 时 argmax 选的深档);③ planner 按真实 bs 定价、执行按 padded_bs 取整(量化未定价);④ tier 边缘按 bs128 推导,不随 bs 迁移。按表内 shelf 重算,即便 tiers 换 [1,2,5],bs64 poetry 的理论净收益也只有 ~+3%。 + +**结论**:B300 agg 大并发下动态裁剪无收益不是 bug,是**成本结构性的**:不可裁地板占一半以上 + θ 宽平台 + 桶阶梯,使"省下的 token"折不成时间;而任何裁剪都真金白银付 accept。收益 regime 在低并发/延迟场景(SGLang 原生 bs1 +17.8% 亦证)与 θ 更陡的硬件上。 + +### 【08-08 23:00】跨栈对照(SGLang × 我们的负载):**不复现我们的无收益——结论修正** + +用户指令:让 SGLang 跑与我们完全相同的负载,表用他们自己的采集方法。设置:他们的 #28 镜像 + 他们的最优配置(Flash tp4/dp4、dsv4 后端、禁 autotune;Pro 在其镜像图捕获阶段崩溃 `Inplace update to inference tensor`,栈级兼容性 bug,已记录)+ 他们的 profiler 现采表(17 探针自检通过)+ **我们的 rounds 负载**(poetry/arena × 512/1024 global × warm+5 reps,同一客户端)。 + +| cell | notrim | compact(表)| delta | +|---|---|---|---| +| poetry 512 | 8744 | 9419 | **+7.7%** | +| poetry 1024 | 6293 | 6455 | +2.6% | +| arena 512 | 7959 | 7933 | −0.3% | +| arena 1024 | 6307 | 6493 | +2.9% | + +- **机制数据(poetry512,decode 日志逐步统计)**:notrim accept_len 2.79 @ cap 6.00;compact accept_len 2.52 @ cap 3.58。即:裁掉 ~40% verify token,付出 −9.7% tokens/step,**净 +7.7%**——他们把省下的 token 兑换成了时间。 +- **与我们的同工况对比**:我们在 128/rank(poetry1024)裁剪深度 (33%) 与 accept 代价 (~−10%) 几乎相同,但净 −1.6%。**两边的"贸易条款"一样,差在结算汇率**——我们的省 token 没有变成省时间。 +- **结论修正**:昨判"B300 结构性无收益"过强。正确表述:**在我们的实现里无收益;SGLang 在同硬件同负载上能兑现 +2.6~+7.7%**(高 accept 的 arena 512 例外 ≈ 0,方向与我们一致)。差距候选(按嫌疑排序): + 1. **图形状经济学**:他们 `align_verify_tokens_to_graph_tier=False` + 每步**均匀 cap**(标量 cap len)——裁剪直接改变每行 token 数,步进更小的图形状;我们是 per-request ragged 窗口 + tier 桶阶梯(#19/#21),裁剪常被桶取整吞掉。 + 2. **每 GPU 负载 regime**:他们 tp4/dp4 = 每 GPU 负载是我们 DEP8 的 2×,θ 份额更陡,兑换率天然更好。**建议补一个对照:我们的栈跑 4 卡 DEP4 @128/rank**,若收益出现则确认 regime 项,若仍无则全归实现项。 + 3. 机制税:他们不裁时 ≈ 0 损耗(arena −0.3%),我们 −3.3%。 +- Caveats:模型不同(Flash accept 2.79 vs Pro 1.45 于 poetry;Pro 在其栈不可跑),per-rank 并发不同(128/256 vs 我们 64/128);+2.6~+2.9% 的两格在 ±3% 噪声边缘,+7.7% 超出噪声带。 +- 后续实验队列:① 我们的栈 DEP4×128/rank 复测;② #19 连续预算(去 tier 量化)原型;③ 均匀 cap 模式(替代 ragged)可行性评估——若三者叠加能兑现 SGLang 级别收益,#25/#30 产品线复活。 + +**【08-09 00:25 追记】regime 判别的结论(实验 ① 改为解析判定)**: +- DEP4-Pro 实测尝试三次均 OOM(Pro 832GB 权重 tp4 = 208GB/卡,图捕获期差 1-2GB;KV 0.32 + 捕获集裁到 5 尺寸仍不够)——Pro 在 4 卡上不可行,与 SGLang 侧 Pro 崩溃互为印证:**跨栈同模型比较做不到,模型差异是本对照的固有 caveat**。 +- 但 regime 问题不需要 DEP4:**per-GPU 负载对齐的格子已经存在**。他们 poetry512(tp4/dp4:每 GPU 128 请求、~128k token MoE 负载)↔ 我们 poetry1024(DEP8:每 GPU 同样 128 请求、~128k token)。同等每 GPU 压力:**他们 +7.7% vs 我们 −1.6%**。→ **差距主体是实现侧**(图形状经济学 + 机制税),不是负载 regime;残余混杂仅剩 EP 宽度(EP8 vs EP4 的 a2a)与模型(Flash vs Pro 的 accept 分布)。 +- 优先级更新:#19 连续预算原型 + 均匀 cap 模式评估升为首位修复路线;DEP4 复测取消。 + +### 【08-09 01:30】SGLang 深挖:兑换率数学闭合,第一嫌疑修正为**批碎片化** + +双方各自的成本表 × 各自实测的裁剪深度/accept 代价: + +| | 步时 | accept | 预测净 | 实测 | +|---|---|---|---|---| +| SGLang poetry512(128/rank 真实步)| 44.9→35.4ms(−21%)| −9.7% | +14.7% | +7.7% | +| 我们 @128 行/步(假想满行)| 92.2→77.2ms(−16%)| −9.8% | **+7.7%** | — | +| 我们 @64 行/步(实际形态)| 60.3→55.3ms(−8%)| −9.8% | **−1.6%** | **−1.6%** | + +- **我们自己的 v2b 表预测:真 128 行/步下裁剪净赚 +7.7%**(与 SGLang 实测同值)。但 poetry1024 的 padded_bs 直方图 = {64: 2364, 72: 443, **128: 仅 134**}/4448 步——**执行器在 1024-global 下多数步只装 64 行**;64 行 regime 地板占 60%,预测 −1.6% = 实测 −1.6%,逐位闭合。 +- **结论**:我们的裁剪机制、决策、成本表全都没大错——**是执行器的批装配把每一步滑进了裁剪必亏的 regime**。SGLang 同压力下每步稳定 126-170 行(decode 日志)。修好碎片化,poetry1024 按自家表应翻到 ~+7.7%,毋须改裁剪本身。 +- **新任务(第一优先)**:定位 1024-global 下 rank 步只有 ~64 行的机制——候选:serve 准入节流、调度器容量交互、DP 路由不均、ragged fit 大批回退。判据:修复后 padded_bs 直方图模式移到 128,poetry1024 转正。 +- 次级(降级但成立):64/rank 低压场景救不回(地板 60%,SGLang 也是靠 128/rank 赚);我们 θ 有平台他们没有(tp4 大形状);地板 29.5 vs 13ms(层数归一后仍 ~1.6×);机制税。SGLang 画像:高 accept 的 arena 他们也只打平(cap 4.01,−0.3%/+2.9%);超出 graph 覆盖(>192/rank)收益缩水到 +2.6%。 + +### 【08-09 01:35】cap-accept 判决(#20):被裁 token 只有 6.2%/16.1% 本会被接受——决策层无罪 + +cap-accept 模式(同一 planner 决策,提交完整块全量验证、接受端封顶,逐请求精确计数反事实)跑同负载: + +| | poetry | arena | +|---|---|---| +| 被裁位置数/请求/步 | 2.08 | 0.71 | +| 其中本会被接受 | **0.130(6.2%)** | **0.114(16.1%)** | +| 占无裁剪接受量 | 9.8% | 4.3% | +| trim_regret_rate | 22.7% | 34.3% | + +- **决策层定案无罪**:poetry 被裁位置 93.8% 反正会被拒,置信筛选有效;arena 上 planner 少裁(0.71 位)且被裁位置存活率高 2.6 倍——行为与风险结构一致。~10% 的 accept 代价与 compact 臂吞吐分解、SGLang 对照(9.7%)三方吻合。 +- **责任链闭合**:决策层 ✅(本节)→ 经济层 ✅(v2b 表:真 128 行/步预测 +7.7%)→ **执行层 ❌(批碎片化,64 行/步 → 地板 60% → −1.6%)**。修复目标唯一:批装配。 +- 更正记录:SGLang profiler 支持双模式(`--fracs` 给定 → 2D `SpsAdditiveCostTable`,离对角 frac 扫描,与我们 v2b 同构且为 #26 的方法来源;省略 → 对角 static → 1D `SpsCostTable`)。#28 与本次跨栈对照用的均为一维模式(manifest fracs=null)——他们拿一维粗表仍 +7.7%,进一步支持"表精度非胜负手"。 + +### 【08-09 03:00-08:00】同台对照终判(parity):裁剪收益 = 慢执行的补偿;我们的栈无肉可裁 + +用户指令的实验:**同模型(Flash-DSpark)、同 4 卡 B300、同 per-rank 负载、同客户端,各自用各自的 profiler 采表**,对照 lbs128/256 下两栈行为。四节点七臂(双方 notrim/trim/cap-accept + 双表)。 + +**① 绝对吞吐(同一切条件,notrim)**: + +| cell | TRT-LLM | SGLang | 比值 | +|---|---|---|---| +| poetry 512 | 20,840 | 8,794 | 2.4× | +| poetry 1024 (lbs256) | 21,838 | 6,324 | **3.5×** | +| arena 512 | 26,279 | 7,941 | 3.3× | +| arena 1024 | 20,041 | 6,397 | 3.1× | + +**② 裁剪收益**:SGLang compact +5.9%(poetry512)/ +0.2~+0.8%(其余);TRT:**profiler 判自家新表 inert(argmax 在 bs32-256×接受率 0.3-0.95 全网格上零次选择裁剪)拒绝落盘**;强制按 SGLang 深度(frac 0.4)裁 → **−12~−22% 全线亏损**,表的判决被实证。 + +**③ 成本结构(双方 profiler 同台采集,机理定案)**:TRT 步时 ≈ 行数的函数——bs128 时 L=2 与 L=5 步时同为 39.2ms(token 翻倍零增长),fixed 20.1ms + α∝行数,**θ≈0,地板 93-95%**;SGLang 步时 ≈ token 的函数(13→45ms 随 M 涨 3.5×,θ 份额 ~70%)。**同一块 GPU、同一个模型,成本结构完全相反——动态裁剪是 token 成本主导执行的补偿机制,在行数成本主导的执行上无标的**。 + +**④ 决策/排序质量(cap-accept 对称测量)**:被裁位置存活率 TRT 15.2%(1.65 位/请求,裸 sigmoid,125 万请求)≈ SGLang 11.4-14.5%(差分估计)≈ Pro+STS 6.2-16.1%(昨日精确计数)——**双方筛子同级,排序不是胜负手**。 + +**⑤ 工程包线**:SGLang decode 图上限实测 192 行(256 行 cublas 捕获崩溃、profiler 自我拒绝),Pro 模型在其栈图捕获崩溃;其每 rank 在飞请求恒饱和 ~110-115(与需求无关)。TRT 256 行图一次捕获成功,notrim 吞吐 512→1024 global 仍在上升。 + +**⑥ 附带修正**:SGLang profiler 亦有二维 additive 表能力(`--fracs` 离对角扫描,#26 方法之源),canonical/blog 流程用一维 static 模式(两次对照均如此);其 STS 机制存在但 canonical 不启用(本对照两侧均未用)。 + +**产品线结论**:B300 + 高效执行(本仓)+ agg 场景下,**动态 verify 裁剪没有经济基础**——不是决策不准(④)、不是表不准(③已双向验证)、而是没有可回收的时间。#19(连续预算)在 θ≈0 下亦无标的。**保留价值**:①(2.4-3.5×)与 lbs256 扩展性本身;#31(Pro@DEP8 批碎片化)仍值得修——Pro 的 θ 份额(40-53%)非零,v2b 预测真 128 行步 +7.7% 未被 Flash 判决否定,须修完碎片化后用 Pro 复测定夺;#25 device 序幕的机制正确性已证(census/GSM8K),其激活条件同上。 + +### 【08-09 08:50-10:15】iter-log 地面真值:θ 平坦结论被推翻,真凶 = 短窗桶 kernel 病灶(#32) + +用户质疑 profiler 可信度(stats 开销污染?pin 是否生效?)→ 重构为纯 `print_iter_log` 的五阶段隔离矩阵(Flash DEP4,逐迭代 `host_step_time`/`prev_device_step_time`/`num_generation_tokens` 地面真值,单 serve 上 pin 换档零重启): + +| 阶段(同负载)| tok/行 | ~128 行 device ms | ~256 行 device ms | +|---|---|---|---| +| U6 notrim(满块均匀图)| 6.30 | 46.4 | 72.3 | +| R5 pin 满窗(ragged 机制,同 token)| 6.30 | 47.9 | 79.8 | +| R5+`enable_iter_perf_stats` | 6.30 | 48.0 | 81.6 | +| **R3 pin 窗口4(token −33%)** | 4.20 | **53.7** | **107.0** | +| R2 pin 窗口3(token −50%)| 3.15 | 50.3 | 102.8 | +| FRAC 0.4(自然抖动)| 4.20 | 53.7 | 102.4 | + +**分解全部闭合**:① pin 真实生效(token 计数逐档精确变化)→ profiler 钉窗机制无罪;② `enable_iter_perf_stats` 税 0-2% → 污染假设基本排除(但采表默认关掉它,已列方法论改进);③ ragged 机制税小(+1.5ms@128 行 / +7.5ms@256 行);④ 桶抖动 ≈ 0(FRAC ≡ pin4);⑤ **真凶:短窗桶形状本身**——token 砍 1/3,步时反涨 **+12%/+34%**,且非单调(窗口 4 比 3 更糟)。**θ 在 serving 真值下不是 ~0,是负的。** + +- **跨证据链**:GB300 Pro frac 表的 w=4 异常(bs64: 62.7 > w5 的 60.8)、我方 Flash profiler 的 42.09 离群点(bs128 L=4)、SGLang #28 bs128 裁剪后 iter 28→33ms——同一病灶跨模型、跨节点、跨栈复现。嫌疑:attention/indexer/DeepGEMM 在非满宽 token 形状下的 tile/算法选型退化。已立 **#32**,下一步 nsys 对拍 pin5 vs pin3 单步 kernel 时间线点名退化 kernel。 +- **对既有结论的修订**:昨夜"θ≈0 → 裁剪无标的"修正为"**满窗形状高度优化、短窗形状病态劣化 → 当前裁剪不但无收益还倒贴**";profiler 引擎模式的平坦行是合成负载(input_len 1024)遮蔽了病灶,采表流程需改用真实负载形态或 serving 日志建表(`--from-iter-log`)。 +- **产品线含义**:#31(碎片化)+ **#32(短窗 kernel)** 是两把钥匙——#32 修复后 θ 恢复正斜率,Pro 的悬崖交换率(GB300 表:w5→w3 省 34% 步时/40% token)将**优于**当前一切估计;裁剪经济学的最终判决推迟到两者修复后的 Pro 复测。 + +### 备选修复方案评估(供评审;当前实施的是 A,标记为暂定) + +- **A(已实施,Python-only)**:捕获拷贝换 D2D 源(`_seq_lens_cuda`)+ NEXT_N 抬全局上界。最小爆炸半径,当场可在活挂载容器里验证;把双胞胎缓冲降级为 canonical 缓冲的衍生物。host 路径装填时刻数值逐位等价,uniform/eager 分支未触碰。 +- **B(否)**:序幕直写该缓冲 + 把 H2D 挪出捕获区。可行但延续"平行缓冲靠写集纪律同步"的模式——本 bug 正是写集漏项造成,管理双胞胎不如消灭双胞胎。 +- **C(架构终态,需 C++ 重编译,建议后续立项)**:改 kernel 删掉 `new_tokens_per_seq` 输入,`nn` 由 kernel 内从 `cu_seq_lens[b+1]−cu_seq_lens[b]`(或 `sp` 直接取图内重算的 `cached_tokens`)导出——冗余输入消失,永不可能再失同步。与 A 不冲突,A 先落地验证机制,C 作 PR 收尾清理候选。 +- **D(否)**:钳制序幕 w ≤ bucket tier。功能倒退且不修主 bug。 +- **E(否)**:仅 device-windows 模式走 D2D。零收益模式分叉,制造未来漂移面。 +- NEXT_N 抬全局上界同时排掉一个 host 侧潜伏雷:#19(取消 tier 阶梯量化)实施后 host 也可能在小 tier bucket 发超 tier 窗口,同样会截断 Phase 1。 + +### #32 nsys 取证:四次报告丢失的根因与 v6 修正(08-09 00:20) + +四轮 nsys 采集(SIGINT 版 / TERM 版 / 手动收尾版 / v5 stop-shutdown 版)全部丢失报告,根因**不是收尾方式,而是采集从未开始**:v5 用 `--capture-range-end=stop-shutdown` 时应用理应在 iter 2050 被自动关停,实测 pin3 应用活过 iter 2110 → `cudaProfilerStart` 触发器从未被 nsys 看见。 + +- **代码侧验证(py_executor.py)**:触发门只有 `iter_counter ∈ profile_start_iters and not is_warmup`(L1739-1748),打印的 `iter = N` 与触发比较用的是**同一个计数器**(L1725),三条执行循环(pp/普通/overlap)都套着 `profile_step`——TRT-LLM 侧路径无问题,且触发时会打 `"Profiling started at iteration N."` 日志(此后可用该行判别触发与注入哪个失效)。 +- **真凶(官方配方对照,perf-analysis.md)**:`trtllm-serve` 的 worker 进程是 **fork-before-exec** 派生的,nsys 缺 `--trace-fork-before-exec=true` 时不会向 fork 出的子进程注入 → worker 里的 `cudaProfilerStart` 对 nsys 不可见,capture 永远不开始,报告永远不落盘。nsys 退出时抱怨的 "re-parented processes" 即其指纹。 +- **v6 修正**:`--trace-fork-before-exec=true` + `--cuda-graph-trace node`(后者让 CUDA graph 内的 kernel 逐个展开——我们全部计算都在图里,做 pin5 vs pin3 kernel 对拍必须要 node 粒度)。两轮已重发(pin5@743961 / pin3@743962)。 + +**v6 判决(08-09 22:10)**:fork 追踪使触发器**开火了**(两轮 serve log 均出现 `"Profiling started at iteration 2000."`,这是 TRT-LLM 侧打的,证明 env 到位、门条件满足),但 nsys 退出时自供 `Processing events... → No reports were generated`——**容器内 nsys 2026.2.1 仍未注入到 worker 进程**,cudaProfilerStart 调了但没人接。此前"没看到 Profiling stopped"是预期行为(该日志只在 torch trace 开启时打,L1666-1671)。另收获:stop-shutdown 关不掉 re-parented worker,`--wait=primary` 让 nsys 干等,需手动 TERM app(留 nsys)才收尾。 + +**v7(执行中)——采用同事已验证配方**:peihengh 的 dsv4 nsys 战役(同集群同拓扑 B300 + trtllm-llmapi-launch + trtllm-serve)成功出过 `.nsys-rep`,其 start_worker.sh 注释直言 *"Lingjie's hang-avoidance: NSYS_BIN=extracted 2025.5.1"*——与用户"nsys 得用 2025.5.x"的提示互相印证:**容器自带 2026.2.1 有问题,须用 lustre 上提取版 2025.5.1.121**。v7 变更:① nsys 换 `peihengh/dsv4_pipe_clean_20260723/tools/nsys-2025.5.1.121`(容器内烟测可运行);② `--capture-range-end=stop`(cudaProfilerStop 时就地落盘,app 继续跑,绕开 shutdown 死结);③ 照抄其余旗标(`-t cuda,nvtx --cuda-graph-trace node --gpu-metrics-devices=none`);④ **窗口对准**:pin3 v6 的窗口正好砸进 burst 排水缝(iter 2000-2025 只有 1 个请求;pin5 幸运落在稳态 120 请求)——burst 跨度 ≈ max_tokens/accept 与批大小无关(pin5 ~470/pin3 ~515 iter,确定性强),v7 round 先放校准 burst 实测跨度,再用小 filler burst(64 req × 按缺口折算的 max_tokens)把 2000-2050 顶进整 burst 的稳态中段。 + +### #32 nsys 取证 v7:报告到手 + 初步 kernel 对拍(08-09 22:50) + +**五连丢终结**:v7(提取版 nsys 2025.5.1.121 + `--capture-range-end=stop` + 窗口对准)两轮均落盘 27MB 报告(`nsys/pin5_rank0.nsys-rep` / `pin3_rank0.nsys-rep`),窗口 iter 2000-2050,graph 内 kernel node 粒度。对准偏差:两窗都落在 burst 的 ~65% 位置(排水段),pin5 占用 72→23 请求、pin3 57→6——pin3 没够到 padded_bs=128 档,已再发 v7b 两轮(对准收紧到 burst 20-35% 段)补清洁样本。 + +**混形状粗对拍(初步,待 v7b 清洁窗验证)**: +| 信号 | pin5(6 tok/req) | pin3(4 tok/req) | 读法 | +|---|---|---|---| +| MoE FC1 大 tile `t128x256x256`(swiGlu) | 17.0%,中位 178µs | 3.2%(几乎退场) | 大桶专属,效率最高 | +| MoE FC1 中 tile `t128x128x256u2` | 1.7% | **8.0%,中位 110µs** | 小桶主力:token 减半,GEMM 只省 38% | +| MoE FC1 小 tile `t128x64x256u2` | 0.1%,中位 69µs | **5.7%,中位 94µs(m3/m5=1.35)** | 同名 kernel 在 pin3 上下文里慢 35% | +| FMHA sparse / A2A / topK / 元素级 | — | 比值 0.6-0.8 | 正常随 token 缩放 | + +**初步指认**:短窗桶的病灶集中在 **MoE grouped GEMM 的 tile 阶梯**——token 桶变小 → 换小 tile → 每 token 效率近乎腰斩,GEMM 时间不随 token 线性下降,形成 #32 的"常数惩罚";与 GB300 Pro frac 表的平台+悬崖楼梯(MoE 更重 → 惩罚更大)跨模型同构。另见 pin3 独有的 FMHA 变体(1.0%)待查。**定论等 v7b 的 (128,768) vs (128,384) 同批次干净对拍。** + +### #32 定论:短窗惩罚 = `_prepare_inputs` 的 host 同步链,不是 kernel 病态(08-09 23:40,v7b 清洁窗) + +v7b 两轮对准命中:pin5b(122±3 请求,桶 768)/ pin3b(122±2 请求,桶 512),窗口 iter 2000-2050 全程均质满批,**同 padded_bs=128 的干净对拍**。步时(窗内 `prev_device_step_time` 均值):pin5b **51.39ms**,pin3b **58.25ms**——token −33%,步时 **+13.3%**,在 nsys 在场下复现 iter-log 矩阵的短窗惩罚。 + +**三层分解(每 50-iter 窗口,rank0):** + +| 层 | pin5b | pin3b | 判读 | +|---|---|---|---| +| GPU kernel 总时间 | 2638.7ms | 2281.9ms(0.865) | kernel 变**快**了,不是 kernel 病态 | +| GPU 空转间隙(>merge)| 162ms(~1 个大隙/步)| **844ms(~7 个/步,单隙 6-8.7ms)** | 慢的全部来自空隙 | +| `_prepare_inputs` NVTX 相位 | 1253.7ms(25.1ms/步)| **1977.4ms(39.5ms/步)** | 惩罚源头相位 | +| 其中 cudaEventSynchronize | 100 次 / 0.9ms(免费)| **300 次 / 1078.2ms(21.6ms/步)** | **6 次/步阻塞 D2H,平均 3.6ms/次** | + +**机理**:短窗(w块 128 → 每 gen 请求锁 2(偶 3)块(_life_cycle_registry.py:38-51)。**容量 = 1.05·maxbs/2.07 ≈ 0.51·maxbs**:maxbs=128 预测 63-65/实测 61-63;maxbs=256 预测 131-135/实测 131-136。配额不敏感(下限无配额项)、kv_util 0.014 是 token 口径(窗口池实为 100% 锁死)、首波满批(短 prompt 1 块,h>128 才要第二块)、spec 降级 vs no-spec 死锁(接受方差去相位 vs +1/步全队锁相集体撞边界)——全部观测闭合。字节估算器还漏了 +draft 窗口放大(cache_manager.py:1316-1326),但字节不是绑定资源。 + +**C 方案实验判决(08-10 20:45,判别象限全齐)**: + +| 象限(Pro@DEP8,poetry 1024 burst)| 每步调度 | 结局 | +|---|---|---| +| spec, maxbs=128 | 60-63 | 优雅降级 | +| spec, maxbs=256 | 115-128 | 健康 | +| no-spec, maxbs=128 | 61 | **死锁崩溃** | +| no-spec, maxbs=256 | 128×254 iter | 健康 | +| **no-spec, maxbs=128 + `max_seq_len: 8192`** | **128×230 iter** | **健康,零死锁** | + +`max_seq_len` 显式设为真实值即解钉(推断值 1M 经 typical_seq_len 把窗口池饿到下限)——**C 方案(部署侧)验证通过,按用户决定采用**。部署规则:DSv4 serving 必须显式设 `max_seq_len`(或 `avg_seq_len`),且 `max_batch_size ≥ 目标并发/rank`(设了 max_seq_len 后无需 2×)。v2 定容的代码级修复(min-slots 稳态建模 / typical_seq_len 默认值 / 池钉死告警)属上游 KV-v2 关注面,与本 PR 分离,另开 PR 跟进。 + +**修复**:A(首选)min-slots 按稳态 decode(2-3 块/请求)建模;B 估算器补 draft(须配 A);C `typical_seq_len` 不默认 max_seq_len / 部署侧显式设 max_seq_len 解钉(判别实验在跑)。上游三报:下限建模、估算器漏项、no-spec 高并发死锁崩溃(报错信息误导)。附:python-scheduler 旗标在 v2 下是 no-op(_util.py:2618),DSv4 强制 v2 管理器——当晚两条二分腿均为无效对照,已修订。 + +### #34 实验环闭合:2×2 象限 + C 方案判决(08-10 21:05) + +| 象限(Pro,poetry 1024 burst) | 结果 | +|---|---| +| spec, maxbs=128 | 稳态 ~60-63(优雅降级) | +| spec, maxbs=256 | 满批 115-128 | +| **no-spec, maxbs=128** | **61 处死锁崩溃**(scheduler_v2.py:439) | +| **no-spec, maxbs=256** | **满批 128,零死锁**(254 iter) | +| **no-spec, maxbs=128 + `max_seq_len: 8192`** | **满批 127-128,零死锁(254 iter)——C 方案判决命中** | + +结论:上限=窗口池按 min_slots 下限定容(比例定容被推断的 max_seq_len=1M 饿死)+ 运行时每请求锁 2 块;**显式设 max_seq_len 即解钉**(用户已拍板 C 方向)。部署规约:DSv4 serving 必须显式设 `max_seq_len`(或 `avg_seq_len`),否则窗口池被钉死在 ~maxbs/2 并发,no-spec 高并发直接死锁。代码侧 C 修复(typical_seq_len 不默认 max_seq_len)+ 上游三报(下限建模、估算器漏 draft、死锁而非降级)待落。 + +### GSM8K 精度回归:新表真实裁剪下通过(08-10 21:40) + +`test_gsm8k_dep8_ragged_verify[moe_backend=TRTLLM]`(DEP8 + overlap + CUDA graph + ragged verify + 置信度调度),`TLLM_DSPARK_SPS_TABLE=postfix_pro_table.json` 注入修复后真实表(planner 主动裁剪而非退化均匀),**1 passed(27 分钟)——精度 ≥ 记录参考值,断言通过**。终局判决的精度一环闭合。注:MEGAMOE_DEEPGEMM 变体在该节点因 DeepGEMM JIT "NVCC compilation failed" 无法初始化(nvcc 在位、/tmp 空,原因未深究;TRTLLM 变体即吞吐判决所用后端,更具代表性)。 + +### Disagg 战役:环境坑编年史(08-11 06:00,第 4 轮进行中) + +用户转向 disagg 形态(只测 gen server 吞吐;纯 gen 步 = 裁剪最有利形态,并绕开 agg 的 ctx 准入带宽上限)。协议:每 (臂,local bs) cell = 独立节点对 + 全新 ctx/gen/router,gen maxbs=当档,客户端超发 16×LBS 钉满,3×2 矩阵({32,64,128},256 搁置)。 + +**环境坑矩阵(全部非本 PR 代码问题):** + +| 轮 | 配置 | 结局 | 教训 | +|---|---|---|---| +| 08-07(#29 时代)| NIXL+PYTHON,无 env | gen 死于 #23 修复前的 IMA,传输从未被行使 | 今日 gen 全程健康 → **#23 修复在 disagg 路径成立** | +| R1/R2(今日)| NIXL+PYTHON,无 env | ①router 忘 `--mpi=pmix` 秒死;②修后 ctx 侧 NIXL `TransferState.FAILURE` → gen 等 KV 超时自杀 | 任何 trtllm CLI 经 srun 都要 pmix(已入记忆) | +| R3 | DEFAULT + UCX 三件套(unset UCX_TLS/关 UCC/RNDV_RAILS=2,抄自同集群日跑 harness)| ctx 启动死于 **v2 KV 管理器 CUDA_ERROR_INVALID_CONTEXT** | DEFAULT(C++ 传输)与 v2 管理器不兼容;p25 时代选 NIXL+**PYTHON runtime** 正是 v2 兼容路径 | +| R4 | **NIXL+PYTHON + UCX 三件套**(未试组合:NIXL 底层即 UCX,三件套正治 UCX 传输)| 六 cell 全过 bringup(证 pmix/PYTHON 组合成立),但传输仍全灭:ctx 首笔 **等 336s 才 FAILURE**(非 0 秒即拒),此后 agent 元数据被 invalidate → 后续全部秒拒 `metadata for remote agent not found`;gen 60s KV 超时先于 336s → 1024 全灭;gen 侧 v2 IMA 为下游殉爆(page 释放撞未完成 NIXL op)| UCX 三件套治不了它——根因在更上层(见 R4 收网)| + +**R4 收网(08-11 08:20,代码取证 + 双实证)——根因:UCX MNNVL 误判 v2 fabric-handle 显存**: +- **机理**:v2 管理器用 `CU_MEM_HANDLE_TYPE_FABRIC` 分配 KV 物理内存(`kv_cache_manager_v2/_cuda_virt_mem.py`);B300 上 UCX 的 cuda_ipc-MNNVL 路径声称此类显存跨节点可达,但两个独立 slurm job 间无 IMEX 通道 → RDMA WRITE 卡死至 UCX 判 ep 死(~336s,恰为观测值)→ NIXL invalidate 对端 → 雪崩。修法即官方文档旋钮:**`UCX_CUDA_IPC_ENABLE_MNNVL=n`**(disagg-serving.md;v1/C++ 栈用普通 cudaMalloc 无此病,故同集群他人 disagg 正常)。 +- **实证互洽**:跨节点 `ucx_perftest`(host 与 cuda 内存、4MB 消息)默认 env 下 90 GB/s 全绿——因 perftest 用普通 cudaMalloc,**对本 bug 是假阴性**;1/1024 偶发成功与 stall(非不可达)一致。 +- **伴随修正**:ctx yaml 原来没给 speculative_config → ctx/gen 窗口池几何不一致(128/8 vs 133/13)且 draft KV 无人 prefill;已镜像 DSpark spec 块到 ctx(MTP disagg 同款做法)。取证副产品:两次 transceiver init/侧 = KV 估算临时 executor(正常);`kv_transfer_timeout_ms` 默认 60s 且**传输无重试**;PYTHON runtime 为 v2 唯一兼容路径(官方注释坐实)。R4 现场已归档 `dspark-runs/r4_archive/`。 +- **R5(进行中)**:MNNVL=n + ctx spec parity,金丝雀 S:128 先行,通过后其余五 cell 并行重跑。 + +### 对齐实验:zhaoyuanh 的 throughput_1k 负载 × 我们的栈(08-11 02:10,主矩阵完成) + +**协议**(与 roofline 报告对齐,按用户修正):throughput_1k parquet 前 1024 行 file 序(源 parquet,SHA 与其冻结 dump 不同,已标注)、greedy temp=0、max_tokens=64、**持续满并发**(完成即补位,计量段=256 个 ramp 完成后的 1024 个完成);**每 cell 独立重启且 maxbs=local bs**;max_seq_len=8192 显式;成本表与 STS **均在本负载重采**(表:burst 模式采纯 gen 段,bs32/64 实测、128/256 α 外推;STS:142k 样本,位置接受率 [0.70,0.46,0.30,0.20,0.15],温度 ~2 → 裸分数显著过自信,ECE 5/5 位置改善)。 + +**主矩阵(中位数 tok/s,n=3/cell)**: + +| local bs | notrim | sched(表+STS) | Δ | pin3 参照 | +|---|---|---|---|---| +| 32 | 1,850 | 1,734 | −6.3% | — | +| 64 | 2,988 | 2,858 | −4.4% | — | +| 128 | 4,336 | 4,284 | −1.2% | 4,227(−2.5%)| +| 256 | 7,216 | 7,112 | −1.4% | 6,129(**−15.1%**)| + +**判读**:①此负载上调度裁剪无收益(−1~−6%),盲砍(pin3)在大 bs 重亏——**与 poetry/arena 的 +5.9~+24.7% 并立:裁剪收益是负载性质的函数**;②机理:64 输出 + 长 prompt + 持续补位 → 几乎每步都是 prefill 主导的混合步,纯 gen 步采的成本表在混合步里**高估 θ 节省**,被裁的接受 token 却真实损失;③他们 +13.29% 的天花板在我们栈的混合步现实中不兑现,落在其自身敏感性分析预言的失效区(静态代理 + 紧凑执行假设);其"DL3 优于 DL5"在我们 runtime 亦反向(诚实支付 D(5) 的 pin3 −15%);④**产品方向**:planner 需要"无标的自保"——混合步密集时降权或停裁(表纳入 ctx 混合度维度)。no-STS 消融进行中(分离校准效应)。 + +### 产品线终局判决:Pro@DEP8 置信度调度全面转正(08-10 13:50) + +Flash 三臂翻案后,Pro@DEP8 终局(maxbs=256 解锁满批 + `postfix_pro_table.json` 新表,scheduled vs notrim,双节点臂序对调,每 cell n=10): + +| 载荷 | 全局 bs | notrim | **scheduled** | Δ | +|---|---|---|---|---| +| poetry | 512 | 12,351 | 13,568 | **+9.9%** | +| poetry | 1024 | 15,586 | 19,431 | **+24.7%** | +| arena | 512 | 15,362 | 16,400 | **+6.8%** | +| arena | 1024 | 20,887 | 22,124 | **+5.9%** | + +**判决**:①四 cell 全正,最难的 poetry@1024 达 **+24.7%**,超过 SGLang blog 量级(H200 +10-20%);②master arc"为什么没有收益"终章——收益从未缺席,被我们自己的三个问题吃掉:#32 prep 同步链(每步 ~14ms)、#31 半批上限(cap=maxbs/2,兼致 2.7× 步时惩罚)、以及嵌入这两个伪影的旧成本表;③修复组合 = b87ea28bf3(双缓冲)+ maxbs=2× 目标并发(部署参数)+ 修复后重采表。 +**方法论注记**:745396 首次 S 臂遇 nvlink peer 显存错误(warmup 期),同配置两节点随后均干净复跑 → 判硬件抖动,非确定性 bug;若复发再立案。 +**余留**:#31 代码级根因(spec 请求 2× 槽记账,已有部署级规避)、GSM8K 精度例行回归、#19 连续预算(θ 单调后价值上升)、反馈 SGLang(2026.2.1 nsys、192 行图上限)。 + +### #31 碎片化:复现 + 排除链 + 工作假设(08-10 05:50,进行中) + +**复现(Pro@DEP8 单 8-GPU 节点,poetry 1024 全局,notrim + iter log)**:rank0 每步调度请求数众数 **60-62**(上尾 75,偶发 65 个 iter@128),远低于公平份额 128;累计分配 639/5120 ≈ 完全均衡 → **路由公平,并发准入被卡 ~60/rank**。一个 1024-burst 要 ~4480 iter 才消化完(等效并发 ~60)。 + +**排除链**:① KV 容量——`free_gpu_memory_fraction` 0.5→0.85 后最终配额 35.4→60.1GiB(定容管线无恙;注意 v2 有两条 quota 日志,第一条是估算期临时缓存,勿误读),直方图不变;KV 需求(~830 token/req × 8438 B/token)远低于配额,排除。② 客户端连接池——nsys_burst2 的 `TCPConnector(limit=600)` 曾是干扰源(600/8=75 恰是上尾!),提到 1300 后众数仍 60-62,排除(**但注意:此前所有用该客户端打 >600 全局并发的实验都被客户端限流污染,幸而 p25fix_bench 是 1100,bs1024 吞吐数据勉强幸免**)。③ ADP per-rank cap(nvbug-6133201 派生上限)——serve log 无 capping 警告,4096/6=682 不咬,排除。 + +**工作假设(判别中)**:调度器步 token 预算按 **KV block 粒度**给 gen 请求计费——Pro block=64 → 4096/64 = **64 请求/步**(众数 60-62 恰在其下,且与最初"64 行步"观察逐位吻合);Flash block=32 → 4096/32 = 128(DEP4 从未撞上,v7b 实测 122 通行)。判别:Pro `max_num_tokens` 4096→8192,若调度跳 ~128 即定案;修复方向 = gen 请求按实际步 token(6)计费,或部署侧提高 max_num_tokens。 + +**⚠ 实验作废公告(08-10 06:35)**:上述 KV-fraction 排除与 max_num_tokens=8192 判别**全部作废**——744499 是 2 节点作业,而各轮的 serve/客户端/pkill srun(-N1)随机落节点,导致**僵尸 serve 跨轮累积**(清点:bia0038 上 12 个、bia0042 上 2 个 trtllm 进程),burst 常打到旧 serve(带统计客户端坐实:1024 全 OK 而新 serve 日志 0 流量)。可信的仅有首轮复现(60-62 众数,fresh serve)。已全节点清理 + 所有 srun 钉死 `-w bia0038`,判别链从头重跑。教训入册:**多节点分配上做单节点实验,每个 srun 必须显式 -w**。 + +**干净判决链(08-10 08:30,单节点钉死后全部重跑)**: +- 复现坐实:1024 全局真实在飞(带统计客户端 1024 全 OK 且与服务端 200 计数一致),rank0 持有满份额 128 活跃请求(`currank_total_requests = 128/1024`),ctx 在前 3 个 iter 内全部完成(准入一次到位),此后**每步只调度 60-63 个**,kv_cache_util 仅 0.014。 +- 干净排除:① KV 配额(35.4 vs 60.1GiB,直方图不变);② max_num_tokens(4096 vs 8192,不变)→ block 粒度计费假设也随之排除;③ 客户端;④ 准入速率。 +- 新观测:步时 ~290ms @ 62 行(缺健康基线,eager 嫌疑未证);SWA 固定成本 8.1GiB = max_batch_size(128) × 63.6MB/请求,定容公式按满批预留(sizing 侧无恙)。 +- 进行中的二分:v1 KV 管理器(`use_kv_cache_manager_v2: false`)vs python 调度器(`use_python_scheduler: true`)各跑一轮——若 v1 恢复 128 → 病灶在 v2 管理器的调度集成;若 python 调度器恢复 → 病灶在 C++ 容量调度器与 v2 的适配。 + +**二分收网(08-10 09:40)——上限 ≈ max_batch_size/2,ماxbs=256 即刻解锁**: +| 判别 | 结果 | 结论 | +|---|---|---| +| v1 KV 管理器 | 仍 ~60-63 | v2 管理器免罪 | +| python 调度器 | 仍 ~60-63 | C++/python 调度器共同层 | +| 关 overlap | 仍 ~62-65 | overlap 免罪 | +| **max_batch_size 128→256** | **调度 115-128(满份额),步时 290→108ms(2.7×)** | **上限 ≈ maxbs/2 坐实** | + +- **部署级解法(立即可用)**:Pro@DEP8 设 `max_batch_size=256`(或 2× 目标并发/rank)→ 128 行满批 + 步时 2.7× 改善;Pro 复测(trim 产品线最终判决)就此解锁。 +- **代码级根因(待定位)**:两个调度器共享的容量/槽记账把每个 DSpark 请求按 ~2 份计(嫌疑:spec 序列槽为 draft 复占、或 GNE 容量对 spec 请求的双倍预留);`get_max_num_sequences`=pp×bs=128 本身无恙。290ms 步时 = 减半批的伴生病(62 行时慢 2.7×,机理与槽争抢/eager 待nsys确认)。 +- 观测注记:cap 尾部到 67-75(非硬 64)——完成中请求可能释放半槽,支持"槽复占"方向。 + +**修复后 Pro 表(08-10 10:55,maxbs=256 满批 + #32 修复,--from-iter-log,4 pin × 3 burst 双节点)**:`postfix_pro_table.json`,bs128 实测 L2/3/4/5 = **80.8/90.3/102.7/110.3ms**(n=47-81/cell)——**完美单调,w4 异常消失,GB300"平台+悬崖"楼梯消失**(确证均为 #32 prep 串行化 + #31 半批的伪影);固定 39.3ms + α(128)=9.1ms,θ(768)=61.9ms = **步时 56% 可裁**(Flash 为 54%,但 Pro 绝对杠杆 2.4×)。w5→w2 省 26.7% 步时。下一步:Pro@DEP8 三臂终局(scheduled-新表 vs notrim,maxbs=256),节点已排。 + +## 【08-08 10:30】协查 agent(laliao 侧)独立复现成立;变体理论证伪;地址拓扑测绘沉淀 + +独立环境(laliao,agg DEP8 2 节点,自家 gsm8k 表 + tiers[1..5],64 并发 percreq census)四轮实验: + +- **独立复现成立**:`FORCE_BUDGET_FRAC=0.25 + DEVICE_WINDOWS + steep` → **64/64 腐蚀**(median 0.13 chars/tok),与主线症状逐位一致。两个复现前置坑沉淀: + 1. **steep 顶格空转陷阱**:窗口全在 ceiling 时 `give=min(a-lo,hi-b)≡0`,序幕恒等,census 假干净(两轮上当)。**空转哨兵必须随行**:`forced_budget_steps` 计数 + `verify_len_hist` 从 {5:...} 转 {2:...}; + 2. **`FORCE_VERIFY_LEN` 在 device 模式被静默忽略**(`decide_verify_budget` 只消费 `_forced_budget_frac`;pin 只在 `decide_verify_lens` 被读)。device 模式钉工况只能用 frac(p25 之名即 frac 0.25)。 +- **变体理论证伪(census 级)**:`FORCE_GREEDY_KEY=1`(0.7 流量强路由 greedy 图 + FORCE_ARGMAX)照样 **64/64 腐蚀** → greedy 图在真 w≠S 下同病。早前"greedy 净 / advanced 脏"的二分是复现条件不齐(temp0 轮次没凑出 w≠S)的伪信号。 +- **共享 metadata 修复证伪**:`TLLM_DSPARK_SHARE_VARIANT_META=1`(advanced key 复用 greedy key 的 metadata 对象,两图烘焙地址全同)机械上成立(840 对 key `n_addr_diffs=0`、健康路径无损)但 **64/64 仍腐蚀** → 变体地址分叉不是病根。 +- **地址拓扑测绘沉淀**(`TLLM_DSPARK_ADDR_DUMP=1`,warmup 末对每对同形状变体图逐名 diff 全部 metadata tensor 的 data_ptr):48+13 个 tensor 中恰 4 个分叉——`_seq_lens_cuda`、`gen_indexer_kv_lens_cuda_runtime`(on_update 每次重绑新 tensor)、`batch_indices_cuda`、`captured_hidden_states`;外加 advanced 副本 `_force_non_greedy_for_capture` 永久卡 True(引擎对象清了、per-key 副本没清;已加清理)。均降级为卫生项。 +- **进行中(S1 判别)**:`TLLM_DSPARK_HOSTLEN_BOUND=1`——`attn_row_kv_lens_host`(→`TokenMajorGenView.host_past_key_value_lengths`,dsa.py:2123 直喂 trtllm-gen C++ op)在 device 模式下持有 SHAPE split 的 S 布局行长,而"host 读都是 bound"契约对 token-major 行不成立(w 布局下行 extent 可超 S extent)。判别:H2D 拷贝后把该 host 数组覆写为全批真上界(设备侧精确值不动)。净 → host 长度契约面定罪;脏 → 该面排除,与主线"跨行分歧检测器"的首个分歧行窗口类别交叉定位。 +- 与"1-2 个幸存请求"吻合的先验:扁平布局前缀和污染模型——首个规模错误行之后全污染,之前的行幸存。若分歧检测器显示分歧总从固定低行号开始,该模型坐实。 +- 复现资产(laliao 侧,可复用):`tmp/dspark_build/devwin/`(devwin_repro.slurm 多臂 + percreq_census.py + agg_config_dev.yaml),overlay=`overlay_mounts_dev.txt`(28 文件全量覆盖 faf2c60935)。 + +## 【08-08 11:40】独立验证判决(laliao 侧):compressor 根因与修复 D2D 双向确认 + +同一独立 harness(agg DEP8 2 节点、自家表、FRAC=0.25+steep、64 并发 @0.7 census)三角闭合: + +- **pre-fix**:64/64 腐蚀(median 0.13); +- **pre-fix + S1 猜想**(`HOSTLEN_BOUND=1`,token-major host 长度覆写真上界):**仍 64/64 腐蚀**(0.094)——负对照通过,`attn_row_kv_lens_host` 契约面排除,"gen_new_tokens 唯一嫌疑"的排他性在独立环境成立; +- **post-fix(57e9d56e8c)**:**0/64 腐蚀**(61 净,median 3.90),序幕 32 次 apply、w≠S 全程在场(hist {2:1708, 3:366, 4:52, 5:2})。 + +代码级四项核查亦过:单位等价(`gen_token_repeat_list` ≡ staging 时 `_seq_lens_cuda[nc:nc+n]`)、捕获区位置(on_update_kv_lens:932)、序幕写集确无此缓冲(apply_device_ragged_layout 写集核对)、D2D 源为捕获期稳定地址且被序幕改写为 w。 + +**代价审计(修复可发布性)**:向量件零成本(删 pinned 缓冲 + 免每步 host 构造);标量件 NEXT_N→全局 6 进了尺寸计算(`num_total_compressed_tokens`、`kv_comp = torch.empty(...)`、compressed_mask arange),但瞬态/图池水位均 ≤ notrim 满块步已付水位(uniform 步本来就按 6 跑),掩蔽靠逐请求精确向量、正确性自洽;净代价 = 裁剪步的压缩器阶段按满块元数据宽度跑,二阶小头,精确定量归 Step 3 三臂。 + +补充验证在跑:dev@temp0 census、host+HOST_STEEP+frac 回归 census。方案 C(kernel 内导出 nn,消灭冗余输入)同意作为 PR 收尾清理候选。 + +## 【08-08 14:30】跨节点四轮复证(laliao 侧,独立环境):#30 "无收益"结论坐实 + +独立复刻 Step 3 规格(4 个独立 node pair × 三臂冷启动 notrim→host→device;poetry(生成)+arena(500 题循环)× bs{512,1024} global × 3 reps × out768 × temp0.7;**本地重拟表**:frac 扫描 7 断点 + 重建 sts_v1;tiers[2,3,4,5],max_batch 128): + +| cell | h/n 四轮 | d/n 三~四轮 | +|---|---|---| +| poetry bs512 | −3.5/−10.2/−3.5/+0.9 | −0.2/−1.0/−2.2/+2.0 | +| poetry bs1024 | −4.7/−5.5/−1.7/−0.9 | +0.5/−1.8/−1.4 | +| arena bs512 | −13.1/−10.2/−6.7/−9.1 | −8.8/−7.4/−6.6 | +| arena bs1024 | −6.3/−11.1/−4.9/−4.8 | (4h 墙截尾,补充轮在跑) | + +- **d/n:poetry ≈ 0±2%(噪声),arena bs512 ≈ −7.6%(三轮方向一致)——无任何正收益**,与主线终局(device/host −6.6%~+0.3%)方向吻合; +- **h/n 在本环境系统性偏负(−3~−10%)**,比主线的 0±4% 更差——可能因素:本地表拟合质量(7 断点 vs 16)、arena 500 题循环的重复前缀结构、节点差异。方向上仍支持同一结论; +- **判决:B300 agg DEP8 + 校准配置下,host/device 动态裁剪无可复现吞吐收益——跨环境(两套表、两套 harness、两组节点)复证成立。**Step 4 门禁(device>host 复现)正式失效;Step 5 门禁同亦失效。 +- **勘误 + 新 bug(Bug #4?):device 臂 arena bs1024 缺格不是 4h 墙(elapsed 均 ~2.5h),是 5/5 可复现的崩溃**——补充轮(2630554,device 先行)在同一 cell 现场抓获:torch `IndexKernel.cu:111` 设备端断言 "index out of bounds"(index 类算子越界)→ `CUDA_ERROR_ASSERT 710` 全 rank SIGABRT。发生在 cell 爬坡段(混合 ctx+gen 步窗口,iter 日志 num_ctx_requests 2→54;注意设备断言浮出有延迟,故不能据此断定肇事步是混合步还是其前的 gen-only 步)。**判别面:host/notrim 同 cell 四轮全过,poetry bs1024(同 128/rank)device 也过,arena bs512 device 也过——特异组合 = DEVICE_WINDOWS × arena 长 prompt 流量 × bs128/rank**。现场:laliao 侧 `tmp/dspark_build/step3/logs/s3device_2630554_serve.log:55198`(断言首现)。吞吐判决不受影响(该 cell 有主线 device/host −2.1% 佐证);但 device 路径在长 prompt 高 bs 下仍有未清的越界索引,建议与"潜伏 bug ①②"并案(候选:混合步/context 加入步与 device budget 单发状态的交互、_seed_context_windows 与 chunked prefill 的槽位索引)。 + +## 【08-09 02:30】bs 曲线补全(laliao 侧):32→256/rank 四点,d/n 谷型、h/n 单调恶化;崩溃 7/7 复现 + +用户加点指令下把 Step 3 网格扩到两端(同配置,256/rank 端重采表 sps_frac_local_2632577 = bs 轴 [64,128,256] 指纹 256): + +| per-rank bs | poetry h/n | poetry d/n | arena h/n | arena d/n | +|---|---|---|---|---| +| 32(256 global) | **+1.4%** | +0.5% | −1.5% | −0.4% | +| 64 | −4.1% | ≈0 | −9.8% | −7.6% | +| 128 | −3.2% | ≈0 | −6.8% | (崩) | +| 256(2048 global,同节点) | −6.0% | **−1.8%** | −1.1% | (崩) | + +- **h/n 随并发加深单调恶化,32/rank 附近归零/转正**——与 #28 SGLang 原生 B300 曲线(bs1 +17.8% / 8-64 平 / 128 −21% / 256 −4~−8%)同形;B300 的"由负转平"边界在 32-64/rank 之间; +- **深水区 device 亏损显著小于 host**(poetry −1.8% vs −6.0%):lag-0 新鲜排序换回的 accept 在深批下开始盖过序幕固定开销的摊薄——机制方向正确,只是在 B300 上翻不了正; +- **IndexKernel 越界崩溃 7/7 复现**(device × arena × {128,256}/rank;poetry 同 bs 全过):与 bs 无关、与 arena 长 prompt/chunked prefill 工况强绑定;2048 global 下 device arena r1 中途崩(部分吞吐 16.2k 无效)。 +- 工程沉淀:并行单臂(跨节点)与同节点配对的 arena h/n 差了 3.8pp——**arena 的节点噪声远大于 poetry**,该数据集的臂间对比必须同节点。 +- 进行中:cap-accept 测量臂(执行满块、窗口记账,cap_trim_tokens 逐 token 回答"被裁 token 真实存活率")。 + +## 【08-09 03:10】cap-accept 直接测量:被裁 token 真实存活率 11.5%(laliao 侧) + +cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完整收官,回答"被裁的 token 实际有多少会被接受": + +- **计划裁剪 1.89M token,其中 219,990 个(11.5%)在满块执行下实际被接受**——零反事实的逐 token 直数; +- 每请求步平均损失 **0.216 token**;算术闭环:封顶 accept 1.66 + 0.216 = 1.88 ≡ notrim 臂实测(1.80-1.88); +- 损失集中:仅 12% 请求步丢 token(丢者均值 1.8,单步 max 4);删失链自洽:31% 被裁请求前缀活到刀口 → 其中 ~55% 下一 token 真被接受 → 全体 11.5%; +- 与跨臂 A/B 估计(~9%)方向一致、A/B 略偏低。 +- **判词**:调度器 ~89% 的刀砍在死草稿上(排序质量好);无收益的根源是那 11.5% 的 accept 换不回 B300 的 kernel 时间(θ 平台),不是调度器裁错了人——给"成本结构论"补上最后一块直接证据。 + +## 【08-09 04:40】IndexKernel 崩溃根因定案(Bug #4):anchor 停车位在满窗满批步越界一格;"arena 特异"定性作废 + +- **判别链**:devfee 臂(DEVICE_WINDOWS + FORCE_BUDGET_FRAC=1.0,poetry,w≡S 恒等)在首个满批 gen-only 步**秒崩**,同一 IndexKernel index-OOB——推翻"device×arena×高bs"表征(poetry 也崩、w≠S 不必要、非爬坡累积)。 +- **根因(一行公式)**:序幕把 anchor token 停在 `dst = real_draft = real_tokens − n_real`(model_engine 序幕 scatter);`draft_tokens_cuda` 容量 = `max_draft_loop_tokens × batch_size`,**当 padded_bs = max_batch 且所有真实行窗口全满时 real_draft 恰等于容量,越界一格**。八次崩溃全解释:arena 高接受 → planner 常发全满窗口 → 满批步必踩;poetry 真实窗口碎(accept~1.2)几乎从不全满 → 幸存;devfee 恒全满 → 秒崩;128/rank 配置缓冲 640、满 bucket 768 同式。协查报告 S4 曾摸到该停车位但只审了语义撞位(判良性),漏审分配边界。 +- **修复(laliao 本地,一行)**:`draft_tokens_cuda` 容量 +1(停车位合法化为 scratch 尾格;bs×K 步进布局不扰、capture 安全)。验证轮(fee2b)在跑:devfee 转净 = 定案。 +- **教训**:①"数据集特异"表征是接受率经由调度器的间接投影,真正的触发几何要用受控恒等臂(frac=1.0)剥离;② 边界审计要同时问"语义对不对"和"分配够不够"。 + +## 【08-09 04:40】verify-len 评估装置入场费(256/rank,poetry,同节点配对,laliao 侧) + +- **头计算费(nohead→static)**:22747 → 22730 tok/s = **−0.07% ≈ 0**(fused 头前向在图内免费); +- **总入场费(nohead→notrim,frac=1.0 全机器零裁剪)**:22459 → 21859 = **−2.7%**; +- **host 机器组件帐单**(sched-timing,3500 稳态步,256/rank):decide ~450-510μs + allgather ~210μs + fit ~140μs + attach ~240-270μs + token-major 行装配 ~400-440μs ≈ **1.4-1.5ms/步 ≈ 5.7μs/请求**;而 host_step≈device_step(120-290ms,KV 深度双峰)——**host 份额被 overlap 完全遮蔽(<1%)**; +- 推论:2.7% 入场费的大头**不在 host 机器也不在头计算**,在 ragged 图的 token-major kernel 呈现差异(uniform 图 vs token-major 视图的 MLA/索引器路径)+ 图内每 token 机构;device 序幕费(notrim→devfee)待 fee2b 修复验证轮补上。 + +## 【08-09 12:30】序幕 Triton 融合(方案 C)落地:60 发射 → 4 kernel,enqueue 8.8×,parity bit-exact(laliao 侧) + +- **实现**:`dspark_device_select_fused.py`(env 门控 `TLLM_DSPARK_FUSED_PROLOGUE=1`,fallback 原路保留):Kernel A(单 CTA:slot gather→stamp fail-open→温度校准→存活→radix-select 预算 topk→cumsum-cycle fill→行映射,radix 复合键复刻 host tie-breaking 到 bit 级)+ Kernel B(bucket 网格:全部 token 级重写,A 写新窗口独立缓冲、B 读旧 split+新窗口,消除 read-before-write)+ 两个块表 gather。停车位 +1 scratch 不变量保留。 +- **两轮真卡教训**:①容器 Triton 拒收 jit 体内模块级全局/降序 static_range/constexpr 上的 not——全部参数化;②首版 e2e 反降 9%:**Triton 隐式整型值类特化**(每个 int 参数的 ==1/÷16 值类进编译缓存键)+ 形状派生 constexpr,ramp 每个 (padded_bs,bucket) 组合都触发数百 ms JIT 且落在流量中——`do_not_specialize` 全覆盖 + 形状参数改 runtime+mask + warmup 预热钩子(生产路径收敛到 4-5 次编译)。 +- **判决(同节点 A/B,256/rank poetry)**:enqueue mean 2052→**243μs**(尾巴消失,mean≈p50),e2e 22548→22556(打平),r1 无编译坑,parity 15/15,越界 0。 +- **综合口径**:序幕 e2e 代价 = 0~3%(取决于 host 压力:serve 口径曾现形 −3.1%,bench 口径 ≈0,本轮 overlap 藏住)——融合拔掉 host 分量后**对 host 压力不再敏感**。若 device 路径未来在低并发/延迟场景复活,这块砖已备好。 + +## 【08-09 15:10】观测开销排污判决:enable_iter_perf_stats 无罪(≤0.2%)(laliao 侧) + +- **设计**:同配方两次真实流量采集(18 phase:frac{0.2-1.0}×conc{512/1024/2048},poetry,新增 `/dspark/budget_frac` 运行时端点免六次冷启动),唯一变量 = `enable_iter_perf_stats`;判决用**同 phase 步时中位数直比**(深度分布天然配对,免分解、免建表)。 +- **判决**:stats-ON vs OFF 步时 **+0.18%**(18 格,散布 −0.5~+1.6%)——stats 机器(含每步 allgather/D2H)代价在噪声内,**此前全部阶梯/曲线结论无需修正**。客户端吞吐 +2.45% 偏 ON 为短 phase span 噪声(方向反于污染假设)。 +- **方法论坑(已沉淀)**:被动 iter-log 样本直接喂分解拟合器必炸——同 (bs,M) 格步时随 KV 深度 120→290ms 漂移,深度方差全进固定项(签名:fixed=258.8ms、θ(179)=0.7ms)。钉接受的本质是控深度;被动样本要么按 kv_cache_util 分箱,要么只做配对比较。 + +## 【08-09 19:30】短窗病灶三方会师 + pin 阶梯直测复现 + H4 候选机理(laliao 侧) + +- **pin 阶梯直测**(`/dspark/verify_len_pin` 运行时端点逐档钉,真实 poetry serving,~250 稳态步 median/格):64/rank 上 **pin4=59.1ms > pin5=56.9ms**(验更少 token 慢 2.2ms),w3→w4 悬崖 +10.3ms;256/rank 恢复单调(111.5<119.4)——**病灶咬中低 bs,深 bs 消退**。与 θ 表倒挂(62.7>60.8)、真实流量 frac0.65>frac1.0(57.7 vs 54.4ms)三方一致。 +- **H4 候选机理(代码级)**:`_pick_dsl_expand`(dsa.py:211)要求 `expand×atom==next_n` 精确因式分解(atom∈FP8{1,2,3,4});**w4→5 token=质数→唯一解 (5,1)→atom=1、5× 行展开=索引器 5× KV 重读,烘焙进该 bucket 捕获图**。w2/3/5(3=1×3, 4=2×2, 6=2×3)全部可摊销。暗合:此前"表自荐 tiers[1,2,5]={2,3,6} token"恰好全避质数——拟合器已从数据学会绕行。与协作 agent 的 H1/H2/H3 并列待 nsys 点名(pin4/pin5 各 100 稳态 iter 的 cuda_gpu_kern_sum 对拍在采,2025.5.1+cudaProfilerApi 配方)。 +- **修法②(tiers 避质数)首验证伪一半**:host[1,2,5] vs host[2,3,4,5] @poetry 64/rank = **−3.2%**——floor 降到 w1 裁得更深,accept 损失盖过避档收益;正确菜单可能是 [2,3,5](保 floor 只删 w4),推荐器需把 accept 损失定价进去(当前只看步时)。 +- **收益区首个正 cell**:GSM8K × 16/rank × devreal(真实 device 窗口+融合序幕)= **+1.8%**(两 rep 一致,甩开 host 3.4pp)——"高接受负载 × 中低 bs × 新鲜排序"的理论收益区实测点亮;加密扫描(8-64/rank × devreal/devalt)在跑。 + +## 【08-09 21:30】✅ 分歧和解:主线 = DEP4×DSv4-Flash,laliao = DEP8×DSv4-Pro(用户确认) + +- 绝对步时 1.5× 差(满窗 79.8 vs 119.4ms @256 行)与病灶档位差(Flash/DEP4 咬 3-4 token 档、Pro/DEP8 咬 5 token 档)全部由 **模型 × 并行度 → GEMM 形状 → tile/tactic 悬崖换位** 解释;两侧数据各自三源自洽,无人测错。 +- **通用律**:verify-len 阶梯的 kernel 成本地形不可跨 (模型×并行度) 迁移——梯子/表/绕档策略必须逐配置 profile,且在代表性 KV 深度下采。laliao 侧 nsys(pin4/pin5,Pro/DEP8)与主线 nsys v3(Flash/DEP4)各自点名各自真凶,并排即完整故事。 +- Pro/DEP8 侧 H4 候选(w4=5token 质数 → `_pick_dsl_expand` atom=1、5× 行展开)待 nsys 证实;Flash/DEP4 侧真凶待主线 v3。 + +## 【08-09 20:10】⚠️ 短窗病灶的跨 agent 分歧,请求 provenance(laliao → 主线) + +双方一致:非单调、w4 最糟。**分歧:w2/w3 的符号**——laliao 侧两份独立数据(pin 阶梯直测:w2/w3 = 46.4/48.8ms 均快于 w5 的 56.9;真实流量 frac 0.2/0.35/0.5 = 46.3/48.1/48.2 均快于 frac1.0 的 54.4)都显示**只有 w4 档为负**;主线判决 4 称"窗口 4/3 的桶比满窗慢 12-34%"。和解假说:① 均匀钉窗 vs topk **混合窗口组成**(同 bucket key 行组成不同,病灶或在混合行为上);② host pin vs **device 序幕重排后**布局;③ 负载/KV 深度/bs 差异。**请主线补 12-34% 的出处(bs、host/device、数据集、窗口组成、量法)**。在钉死前,交集结论可行动:w4(5 token)档在中低 bs 确定为负,梯子应绕开;分歧决定深裁档(w2/w3)是否可用,直接影响最优梯形。laliao 侧 nsys pin4/pin5(均匀)在采;若假说①成立将补 frac0.65(混合)对拍。 + +## 【08-09 23:30】SGLang 字面复现立项(laliao 侧);请主线共享 #28 环境 + +- 用户指令:按 blog reproduction 章节**用 SGLang 本体**复现(此前只做了设计复刻)。章节要点已抓取:镜像 `lmsysorg/sglang:dev-dspark`(PR #30261 @692c5f7d);frontier 臂 = Flash-DSpark×H200×DP4,one_batch_server bs {1..256}×out1024×temp0.7×frontier_prompt(16 拼接 GSM8K);零开销图 = **Pro-DSpark×B300×TP8×cuda-graph-max-bs 4**;动态/固定切换 = `SGLANG_RAGGED_VERIFY_MODE=compact|static`。 +- 进行中:enroot import dev-dspark 镜像(需验 aarch64);checkpoint `deepseek-ai/DeepSeek-V4-{Flash,Pro}-DSpark` HF 公开未 gated(Flash 74 文件 / Pro 92 文件,login 下载带宽 13MB/s 是瓶颈,将用 HF_TOKEN+分段配方)。 +- **请主线共享 #28 的 SGLang 环境**(镜像/venv 路径、checkpoint 位置、launch 参数)——你们已在 B300 跑通 SGLang 原生 DSpark,复用可省一天;也顺带补 #28 的 provenance。 + +## 【08-10 00:40】nsys 点名(Pro/DEP8 w4 真凶 = H1 GEMM tactic);devalt 复证失败;吞吐收益区关闭(laliao 侧) + +- **nsys pin4 vs pin5**(各 100 稳态 iter,cuda_gpu_kern_sum 对拍):①同批 GEMM tactic 换名——pin5 的 `64x64_2cta_TNT`(3245 次)/`256x128` 在 pin4 消失,被 **splitK 变体族**顶替(`tss_64x32_splitK` 4769 次等,净 +~130ms);②`gate_forward_kernel`(MoE 路由门)同 6100 次变贵 48%(56→83ms)——更小 M 落在更差 tile 点;③pin4 GPU kernel 总时长反而更少(3792 vs 3930ms,MoE 省 85ms)但占空比更低(64% vs 69%)——**"省了活、亏了排程"**:更碎的 splitK kernel 拉高发射/间隙。**判决:H1(tactic 选择)为主犯;H4(索引器质数 atom)未现身 top 差异,降级**。与主线 Flash/DEP4 的 H1 首要嫌疑吻合——两模型同类真凶、不同形状表现。 +- **devalt 复证失败**:两轮四格全负(32/rank:−3.0/−1.8;64/rank:−5.2/−1.7),D1c 首轮 +2.3/+1.0 为基线噪声(nohead 轮间 ±3%)。**B300 吞吐口径的收益区搜索正式关闭**:poetry/GSM8K × 1-256/rank × host/device/梯子变体,无任何跨轮稳健的正收益。存活方向:D2b 深上下文(在跑)与延迟口径(blog 的 B300 Pro 配置 cuda-graph-max-bs 4 本质是 latency 场景,SGLang 字面复现 #32 将验证)。 + +## 【08-10 01:30】收益区搜索终局(laliao 侧):B300 吞吐口径全维度关闭;唯余延迟口径待 #32 + +- **D2b 深上下文(ISL8k,chunked prefill)收官**:32/rank h/n −4.6%、d/n +0.3%;128/rank h/n −4.8%、d/n −5.5%——"θ 份额随 KV 深度涨 → 深上下文是主场"假说**证伪**(裁剪省下的照旧换不回 accept+税)。 +- **搜索矩阵全景(全部同节点差分,<5% 信号均经复证纪律)**:负载 {poetry, GSM8K, arena, mixed, ISL8k} × 并发 {1..256}/rank × 配置 {host, device+融合序幕, tiers[1,2,5]} —— **无任何跨轮稳健正收益**;单轮 +1~2% 信号(GSM8K 16/rank devreal、devalt 32-64/rank)复证全部还原为 ±3-5% 基线噪声。 +- **成因四件套(全部实测定价)**:H1 GEMM tactic 病灶(裁剪后 M 形状触发 splitK 碎战术 + gate kernel +48%,nsys 点名)+ accept 损失(cap-accept:被裁 token 11.5% 存活)+ 桶量化税(请求级 13-16% 落地仅 5.8-7.6%)+ token-major 呈现税(−1.6~−2.7%),叠加 B300 >50% 不可裁地板。 +- **唯一未决口子:延迟口径**(blog 的 B300×Pro cuda-graph-max-bs 4 配置、SGLang 原生 bs1 iter +17.8%)——#32 SGLang 字面复现进行中(镜像导入中,checkpoint 公开)。若延迟口径亦无,"DSpark 动态裁剪在 B300 无收益"即全口径终局。 + +## 【08-10 03:30】SGLang 字面复现收官(#32):blog 原版命令 @GB300,动态几乎全线为负——全口径闭卷(laliao 侧) + +- **配方**:`lmsysorg/sglang:dev-dspark` arm64(63 层 overlay 超内核限,计算节点手工压平 + docker ENV 50 条重注入;镜像/权重下载配方沉淀于 sglang_repro/)+ `deepseek-ai/DeepSeek-V4-Flash-DSpark` 180GB(计算节点 HF 出口 370MB/s,5 分钟)+ blog reproduction 章节逐字命令(DP4、flashinfer_mxfp4、one_batch_server bs{1..256}×out1024×temp0.7×frontier_prompt)。 +- **判决(static vs compact,输出吞吐)**:bs1 **+1.8%**(ITL 2.35→2.31ms,唯一正格);bs8 −12.8%、bs32 −9.4%、bs64 −5.2%、bs128 −6.1%、bs256 −1.6%。acc_length 双臂 4.7-4.9 一致(接受率形态复现)。 +- **结论**:同代码、同命令、同权重、同接受率,只换硬件(H200→GB300),blog 的 +10-20% 收益消失——**硬件地形是第一性原因的最干净证据**;与 TRT-LLM 侧全部结论跨栈闭环。#28 的 bs1 +17.8% 未在字面复现中重现(其 checkpoint/commit provenance 存疑)。 +- **DSpark 动态裁剪 @GB300 全口径闭卷**:吞吐(TRT-LLM,五负载×1-256/rank)与延迟(SGLang 字面,bs1-256)、两套实现——无收益;收益属于 θ 陡峭的 H200/H100 地形。 +- 【08-11 追记】上条"闭卷"已被 #31/#32 修复整体翻案(见主线终局判决);SGLang@GB300 compact 为负仍是其栈自身的事实,但"硬件地形第一性"解读降级为负载/实现因素之后的第三位。 + +## 【08-11 09:00】disagg 独立复现战役立项(laliao 侧):自有 harness 跑通路 + 三负载矩阵;请主线提供 poetry 原文 + +用户令 laliao 侧在 disagg 形态独立复现置信度调度收益,与主线 disagg 战役(R5)并行互证。设计与进展: + +- **通路**:不走主线的手搓 ctx/gen/router 脚本,用我自己 08-10 已跑通的 `bench-trtllm-disagg` harness(agentx gen-only 配方,**NIXL+PYTHON 本来就在用**,单 slurm job 3 节点带 `--segment` 同 NVL 机架分配——按 R4 收网的机理,IMEX 就位,`UCX_CUDA_IPC_ENABLE_MNNVL` 坑对单 job 同 segment 可能天然免疫;bring-up 正在验证)。新模型目录 `deepseek-V4-Pro-nvfp4-dsparkconf`(DSPARKCONF_* 提交时环境变量参数化全臂矩阵:SCHED/TABLE/STS/TIERS/GRAPH_MAX/GEN_ENV),harness 的 CONTAINER_MOUNTS 已修成支持 src:dst 文件绑定(python overlay 28 文件,含 b87ea28bf3)。 +- **协议吸收**:feature doc 的 B' 四坑全部落实——PYTHON+NIXL ✓、MNNVL(视 bring-up 判决)、kv_transfer 600s ✓、**ctx/gen spec parity 已修**(ctx 镜像基础 spec 块,confidence 开关 gen-only);max_seq_len 显式(harness 按 isl+osl 渲染)✓;bs 钉法改用主线定稿(maxbs=目标并发 + 客户端 2-4× 适度超发 + iter log 实测中位为准)。 +- **矩阵**:三负载 {poetry, arena, throughput_1k} × lbs {32,64,128,256} × 臂 {notrim, frac04, scheduled};**SPS 与 STS 分负载重采**(A:KV 深度 ~100-800 vs B:~1110,跨用错价;STS 是数据语义×采样温度的函数,A temp0.7 / B greedy);对表纪律:A 主臂裸 sigmoid(与主线 +5.9~+24.7% 同协议),STS 为 A 的消融格、B 的必挂项。**B 在 disagg 翻正与否 = "混合步成本表 ctx 盲"机理的直接检验**——这是本战役最有信息量的一格。 +- **➡️ 请主线**:提供 poetry 的 **verbatim prompt 原文**(provenance 只有描述:memory/rivers/glass/time、≥8 场景、不复用意象;我 step3 自造的 caravan/steppe 不是同一首,接受率分布不可比,对表 A 必须同文)。最好直接把数据集构造脚本或 jsonl 推上来;arena(arena-hard-v0.1 取 prompt 字段)与 throughput_1k(iputterman/speed parquet)provenance 已够我自建。 + +## 【08-11 09:40】disagg 复现开工(laliao 侧):poetry/arena 线已发射;➡️ 请主线推 throughput_1k 抽取件 + +- poetry 原文与 768-cap/temp0.7 修正已收,verbatim prompt 已铸成 trtllm_custom 数据集(poetry/arena 各 8192 行);客户端(benchmark_serving)已开关化:BENCH_KEEP_EOS=1(自然 EOS)+ BENCH_TEMPERATURE=0.7,对齐 burst 协议。 +- bring-up R1 尸检:非传输非代码——**gen DEP8 每 rank 85GB 权重装载撞 harness 1800s 健康等待超时**(ctx 健康、gen 差几分钟;DEP32 每 rank 21GB 所以旧 agentx 跑通)。wait_server 已抬 3600s。R2(poetry lbs64,ctx-spec-parity + 主线钉法 maxbs=lbs + 2× 超发)已在队。 +- **➡️ 请主线**:`iputterman/speed` 是私有 gated,我侧 401,贵区 dspark-runs 路径在本文件系统亦不可达。请把 throughput_1k 的 **前 1024 条抽取件**(file 序,turns[0],建议直接 trtllm_custom jsonl,~5MB)推到 `docs/dspark_bench_assets/`,负载 B 线即可开工——B 在 disagg 翻正与否是"混合步成本表 ctx 盲"机理的直接检验,是三负载矩阵里信息量最大的一格。 + +## 【08-11 11:30】agg 独立复证判决(laliao 侧):frac04 机械裁在 Pro@DEP8 全线为负;Pro 翻案的承重臂只能是 scheduled + +修复后代码(b87ea28bf3 已 overlay)、maxbs=256+显式 max_seq_len、到达性验证干净(iter-log mode=64/128 满档)、双独立节点对 × n=5: + +| 负载 | 全局 bs | 对A frac04 vs notrim | 对B | +|---|---|---|---| +| poetry | 1024 | −5.0% | −2.6% | +| poetry | 512 | −4.8% | −3.3% | +| arena | 1024 | **−16.9%** | **−14.9%** | +| arena | 512 | **−18.8%** | −15.6% | + +- 与主线不矛盾:frac04 转正(+5.3/+7.6%)是 **Flash DEP4** 的数据;Pro 终局从未报 frac04。本判决补上缺格:**Pro 上机械 40% 裁剪稳定亏钱**(accept 损失定价高),Pro 的收益只能靠"按请求该裁才裁"的 scheduled+修复后表。arena 亏得最狠(−15~19%)与 Flash 侧 arena≈0 的方向一致放大——多样负载对机械裁最不友好。 +- 复证下一步已在跑:我环境重采 postfix Pro 表(pin{2,3,4,5} × conc{512,1024,2048} poetry,--from-iter-log),然后 scheduled vs notrim 终局对表主线 +9.9~+24.7%。 +- disagg 进展:R2 死因=ctx 挂 spec 后 host OOM(tp4 单节点 4 rank 双份权重 staging 撑爆 host RAM;slurm oom_kill 实锤)——**ctx 已改 tp8×2 节点**(与 gen 同规格),R3 在队。坑登记:Pro 这个体量,disagg ctx 至少 tp8(或关 TRTLLM_PINNED_WEIGHT_STAGING 试 tp4,未验)。 + +## 【08-11 14:30】disagg bring-up 门禁全过 + agg 复证首轮的 STS 违规教训(laliao 侧) + +- **disagg R4 门禁全过**(poetry lbs64,ctx tp8×2 节点 + gen DEP8×2 + 两侧显式 max_seq_len 8192):**77,823 步 mode=64 逐字满档**、KV 传输全程行使零 FAILURE、无 IMA/OOM。单 job 同 NVL segment(--segment 分配)下 **MNNVL 免疫假设证实**——五件套里的 MNNVL=n / UCX_TLS 去 IB 是跨 job 跨机架的处方,同机架单 job 不需要。仅墙钟 2.5h 被截尾,已抬 4h。 +- **两个新坑入册**:①ctx(tp4 单节点)挂 spec 后 host OOM(双份权重 staging;Pro 体量 disagg ctx 至少 tp8);②harness 自动推导 ctx `max_seq_len=isl+40`,poetry 短 ISL 下= **104 < 窗口池几何 133**,稀疏 topk 被钳成非 4 倍数 → `SparseAttnTopK must be a multiple of 4` 在 ctx warmup 验证步断言(ctx+spec+短 ISL 的新角落,ISL≈1110 的负载永不触发)。修法都指回协议:**两侧显式 8192**。 +- **我环境的 postfix Pro 表出炉**(pin{2..5}×conc{512,1024,2048} poetry,--from-iter-log):fixed 35.1ms、α(256)=3.8ms、θ 单调(192→7.7 … 768→40.9 … 1536→78.9ms);**bs128 步时 54% 可裁 / w5→w2 省 27%,与主线表(56% / 26.7%)跨环境互证**。login 无 tensorrt_llm 的拟合配方:venv numpy + 包级 stub runner(绕开重 __init__)。 +- **agg scheduled 复证首轮为负(−4~−12%)——但发现协议违规**:我的 config 挂了旧 `sts_v1_reconstructed.json`,主线 Pro 的 +9.9~+24.7% 是**裸 sigmoid**;过时温度扭曲排序即乱裁。v2(去 STS,协议逐字对齐)已重排。⚠️ 若 v2 仍负,则分歧升级为真(候选:表绝对值差异/图集 [1,8,32,64,128] vs 全阶梯/客户端 burst 语义),将请主线共享 `pro_sched_arm.yaml` 逐键对表。 +- stage A 第一波在队:poetry lbs{64,128} × {pin2..5, notrim+STS 采集},disagg 表与 STS 原地重采。 + +## 【08-11 22:30】agg 独立复证闭环:客户端语义定罪 + 统一律三面验证;disagg 发现 gen 不提交进度的疑似分支回归(laliao 侧) + +**agg 复证故事线完整闭环**(全部修复后代码、我自己重采的表、到达性逐字验证): + +| 轮 | 客户端 | scheduled vs notrim | 判读 | +|---|---|---|---| +| schedv2 | 3× 补位闭环(混合步 11%) | poetry −6.3/−6.5%,arena −7.7/−9.8% | 分歧真实但根因在负载形态 | +| **schedv3-burst** | **单波 burst-drain(混合步 ~1%,对齐主线协议)** | **poetry +4.6%/+7.8%**,arena +0.4/−2.4% | **定性复现主线翻案**;唯一变量是客户端 | +| schedv3-gate | 补位 + SKIP_MIXED_TRIM | 门控只回收 +0.5~+2.6pp | 补位损失大头 = 纯 gen 密度稀释(regime 本身),错误裁剪只占小头 | + +- 到达性与混合步密度均 iter-log 直测:两臂 gen 行数众数逐字打满(64/128),补位客户端 11% 的 gen 步带 prefill,burst ~1%。**统一律(收益=纯 gen 密度×bs 深度)在独立环境三面验证:同机同表,换负载形态即可让收益在 +7.8% 与 −10% 间摆动。** +- 遗留:poetry 幅度差(我 +7.8 vs 主线 +24.7 @1024;候选:表绝对步时 1.45× 差(76 vs 110ms bs128 满窗,KV 深度/采集负载不同)、图集 [1..256] 全阶梯 vs [1,8,32,64,128]、burst 计量细节);arena ≈0 vs 主线 +6.8/+5.9(请主线提供 arena 臂的 burst 语义与 OSL cap 细节对表)。 +- frac04 机械裁补格判决(前报)与 schedv3 合并看:**Pro 的收益只有 scheduled+修复后表+纯 gen regime 三者同时成立才兑现**。 + +**disagg 疑似分支回归(阻塞我方 disagg 战役,请主线协查)**: +- 症状:gen server 每步执行真实完整 verify+draft(102ms 设备时间、384 tok/步)但**永不提交进度**——KV util 0.005 平线 129k 步、active 恒 64/512、零完成;1024 请求完整走过 router→ctx→传输→gen,全链无错。pin/sched/FRAC 各臂均复现,与置信度调度开关无关。 +- 反证:我 08-10 旧镜像(b8bc42c3ae,同 harness 同 gen_only 模式,DSpark mtp6)完成流通;主线 Flash disagg cell 亦通。判别局 D1(同 poetry cell × 旧镜像无 overlay)在跑:通则回归定罪(faf2c60935+overlay 代码代的 disagg gen 接受提交路径),再二分 overlay 28 文件 vs sqsh 基底。 +- 附:到达性机制无懈可击(6 个僵尸 pin cell 的 mode 全部逐字满档),disagg 表/STS/三臂在 bug 后面排队。 +- ⚠️ 请主线看一眼:你们 disagg gen 侧用的 commit/镜像与 gen 第一步(消费 ctx 交接后)的提交路径,近期有无已知改动? + +## 【08-12 05:30】量化经济学实锤(fp8 量级轮)+ 冻结 bug 判别链收缩(laliao 侧) + +**fp8 量级复证**(发现我 07-30 已下完 fp8 原版检查点——.incomplete 残留是 hf-download 死垃圾,尺寸校验其实早通过;主线 pro_sched_arm.yaml 逐键 + 他们的表 + burst 客户端,同机同代码只换量化): + +| 负载/bs | nvfp4 | **fp8** | 主线 fp8 | +|---|---|---|---| +| poetry 512 | +4.6% | +3.9% | +9.9% | +| poetry 1024 | +7.8% | **+11.8%** | +24.7% | +| arena 512 | +0.4% | **+5.5%** | +6.8% | +| arena 1024 | −2.4% | +0.3% | +5.9% | + +- **verify 越贵收益越大的量化档位效应直接验证**(poetry@1024 +4pp、arena@512 从零翻正);逐键对表功不可没(头号差异就是 fp8 vs nvfp4 checkpoint + moe TRTLLM vs MEGAMOE_CUTEDSL)。产品含义:**报 DSpark 收益必须标注检查点量化;nvfp4 部署下天花板显著低于 fp8**。 +- 与主线剩余差距(11.8 vs 24.7)第二阶候选:N 臂定义(我 FRAC=1.0+SCHED_TIMING 有机器税,主线是关 confidence 开关——下轮用 CFG_NOHEAD 式 N 臂)、客户端实现(p25fix asyncio vs benchmark_serving)、warm-rep 剔除。 +- 注:本轮 arena 绝对吞吐显著高于 nvfp4 轮(63k vs 40k tok/s @1024)——burst 单波 completion-span 口径对 ISL 分布敏感,横向绝对值勿直接比,臂内 delta 有效。 + +**disagg 冻结 bug 判别链**:D2(去 overlay)= **ZOMBIE(20253 步)→ overlay 28 文件无罪**;D3(ctx 去 spec)= ctx 崩 v2 `dangling rawref` → 该路不通,反证 trap 4,ctx-spec 修复方向无罪;D4(旧镜像 b8bc42c3ae × ctx tp4 无 spec × 冻结组合 poetry/lbs64/benchmark_serving)在队——健康则回归区间锁定 b8bc42c3ae→faf2c60935 两构建之间,冻结则配置组合(客户端/streaming/lbs)定罪。 + +## 【08-12 17:30】AA slo20 战役(laliao 侧):ragged path 生产形态验证通过;冻结案结案=短 ISL;slo20 表 v1 出炉 + +- **冻结案结案**:D7b(主线运行时×fp8×我 cell 仍冻)+ D9(PR HEAD 新编译 sqsh 仍冻)+ AA(39k ISL 同栈同码健康)→ **触发条件 = 短 ISL(64)× DSpark disagg gen**(浅 KV 角落,与 ctx 侧 SparseAttnTopK 短 ISL 断言同族)。修复挂上游清单;所有真实负载(≥1k ISL)不受影响。 +- **AA slo20 生产形态 ragged 验证 ✅**(aa-master 流水线 fork:slo20 keys × nvfp4 DSpark5 × AA 轨迹客户端 × PR HEAD sqsh):58,144 步 97% ragged、**trim_ratio 0.50**、verify_len 全谱 {1:335k, 2:700k, 3:100k, 4:19k, 5:212k}、accept 2.17、graph 零 miss、154 轨迹完成;客户端 p50 146.8 vs notrim 冒烟 131.2 tok/s(+11.9%,rep=1 方向信号,不作数)。 +- **slo20 SPS 表 v1**(pin{2..5}×conc{32,80},--from-iter-log,28 cells):fixed 15.6ms、α(64/rank)=3.9ms、θ 单调(20→4.0 … 256→26.6 … 384→34.3ms)。几何坑入册:**幂次 bs 网格下 pin 比值 {3,4,5} 互不相连,pin5(比值 6)是连接器**——阶梯没有 pin5 拟合必炸 SweepGeometryError。 +- **reuse 纪律**(AA 负载 reuse=true 是产品行为):终局 A/B 每臂独立冷启动 job + settle 240s + 计量 ≥900s 暖稳态、同 seed 对称升温;SPS 采集豁免(按步形状分箱,缓存温度只改箱填充速度)。 +- 流水线坑五连修入册:helpers 全量拷贝、OUT_DIR 自挂载、ctx 模板 MTP spec 块换 DSpark(KeyError 'weight';TRTLLM moe 装 nvfp4 无罪)、表路径需显式挂载(AA 无整卷 /lustre)、STS 采集配置不得带 enable_ragged_verify(验证器要求 ragged 必须有表)。 +- 在飞:STS v2 采集;下一棒 = 拟温度 → 三臂冷启动矩阵(notrim vs scheduled@新表裸sigmoid vs scheduled@新表+STS)。 + +## 【08-12 23:50】AA slo20 终局判决(laliao 侧):scheduled 在生产负载 disagg 上 +6.7%/+7.4%(中位,n=3 冷启动) + +**协议**:slo20 生产形态(2ctx:2gen,maxbs256/mnt1024/131k窗/reuse on/autotuner off)× nvfp4 DSpark5 × AA coding-agent 真实轨迹(32 users,settle+暖稳态计量,seed 42)× PR HEAD sqsh;S 臂 = aa_slo20_table_v1(本形态 pin 阶梯重采)+ tiers[2,3,4,5] + 裸 sigmoid(STS 按用户指令跳过;AA 位置接受率 [0.982→0.892] 极高,裸分数够用);**每臂每 rep 独立冷启动 job**(reuse 纪律:臂间缓存温度对称)。 + +| arm | 生成 p50 (tok/s/user) ×3 | e2e p50 ×3 | +|---|---|---| +| notrim | 130.4 / 130.1 /(r2 NODE_FAIL)| 116.8 / 116.5 | +| sched | **141.7 / 139.0 / 127.2** | **129.8 / 125.3 / 115.5** | + +- **中位 +6.7%(gen)/ +7.4%(e2e)**;trim 0.50、accept 2.21-2.23 三 rep 一致;无腐蚀、轨迹完成流健康。 +- sched r3 与基线打平——rep 间波动 ±5%,n=3 中位口径成立,收窄区间需加 rep;notrim 双 rep 130.4/130.1 极稳。 +- **战役闭环**:ragged path 在 AA 生产形态功能有效(97% ragged 步、trim 50%)且**性能兑现 +7%**——统一律第五格(纯 gen disagg × 高接受产品负载)为正。 +- 过程沉淀(全部入册):AA 流水线十坑(helpers/OUT_DIR 挂载/ctx-MTP-spec/表挂载/STS 五护栏/GEN-only env/…);static×disagg 置信度 NaN(bug 数据点);短 ISL 冻结(已结案);SPS 几何(pin5=连接器);accept 先兆:正式表让 accept 从冒烟 2.17 → 2.22(planner 用真 θ 裁得更聪明)。 +- 余留:sched rep 加密(收窄 ±5%)、c80 档、STS 消融(如果之后想验标定增益)、短 ISL 与 static-NaN 两个上游修复。 + +## 【08-16 02:00】slo20-true 战役(laliao 侧):v2 dangling-rawref 定罪(压力路径);slo20 原件 ×u128 = 99.9% 黄金点;四臂矩阵开跑 + +**真凶(带完整栈)**:压力上来时 gen 多 rank 同步崩:`Error in event loop: Dereferencing a dangling rawref` @ `_apply_disagg_transfer_admission → _revert_deferred_disagg_gen_init_alloc → _revert_ctx_alloc`(py_executor.py:3793)→ executor 死 → uvicorn 僵尸收单 → 客户端 900s 零 chunk 超时 → 会话冷重启死亡螺旋。**只在传输准入的回滚路径触发**(资源压力/deferred 分配),所以 u128 低压 99.9% 绿、u512/c2600 必崩;与 D3 的 ctx 侧 dangling rawref 同族 = v2 生命周期 bug。主线 c2600 不踩的解释:你们的运行时是 main a8b540912c+wheel(新 v2),我们分支自建 sqsh(旧 v2)——**请确认 main 近期是否修过 rawref 生命周期**;分支自建用户都会踩。 +- 连带修正:此前"短 ISL 冻结"定罪降级为存疑(当时判别在 num_instances 带毒渲染器时代;见下)。 +- **自伤教训入册**:fork slo20 launcher 做 ctx reshape(6×tp4→3×tp8,DSpark 双份装载 OOM 所迫)时漏改 `server_config.num_instances` → router 永等 6 实例 → 一整代判别局(ADP/env/容量/swap 全系列)被毒,全部作废重跑。**渲染器改拓扑必须三处联动:启动命令/urls/num_instances。** +- 搬家坑清单(全修):client 需 .venv 直调(镜像无 uv,--quiet 吞错)、base_url 烤死原节点、wait 1800s、DL 臂需显式 block_size=DL。 +- **黄金点:slo20 全套原件(ADP+DEEPGEMM+maxbs256+kv0.8)× u128 = 28,784 成功/27 失败**。与原始 slo20 唯二偏差:ctx 3×tp8、并发 128(均有硬归因)。 +- **四臂矩阵已在该点开跑**:notrim(r1=28.8k ✓)/DL4/DL3/sched(nvfp4 DSpark,sched 挂 aa_slo20_table_v1 + tiers2345 裸 sigmoid)——回答"收益属于置信度调度还是静态砍 draft len"(roofline best-static 之争的实测版)。 + +## 【08-16 09:30】slo20-true 四臂矩阵终表(laliao 侧):静态砍 DL 伤 SLO;sched 无损无益(nvfp4 平 θ 下理性拒裁) + +工作点 = slo20 原件键(ADP+DEEPGEMM+maxbs256+kv0.8)× u128 × nvfp4 DSpark,每臂独立冷启动,AA client 真轨迹: + +| 臂 | 完成数 | SSE p50 tok/s | +|---|---|---| +| notrim(heur)×2 | 28,784 / 29,698 | 92.1 / 91.6 | +| notrim-noheur(公平对照) | **30,199** | 91.7 | +| **sched**(表+tiers2345 裸sigmoid) | 29,179 | **92.5** | +| DL4 | 30,090 | **84.1(−8%)** | +| DL3 | 29,227 | **77.3(−16%)** | + +- **静态砍 DL 在 SLO 语义下净损**:吞吐 +1~4%(rep 方差 ~3% 内)换每请求速度 −8~−16%——roofline 的 best-static 之争在 SLO 判据下出局。 +- **sched 无损无益**:vs 公平对照完成数 −3.4%(方差内)、速度 +0.9%;trim 仅 0.0023、accept 3.55——planner 按表判定 nvfp4 的 θ 不值得裁并拒裁,行为正确。与量化经济学自洽(收益的肉在 fp8,agg 侧 +11.8% 已证)。sched×heuristic_topk 被验证器禁(ragged 批 next_n 重建不成立)——sched 部署需关 heuristic,公平对照已扣除该差异。 +- 下一格候选:fp8 原版 × 同矩阵(唯一可能让 sched 翻正的工作点);u512+ 需先修 v2 rawref(回滚路径)。 +## 【08-18 10:30】slo20 形态成本表重采(laliao 侧):**78% 步时不可裁** —— 拒裁的定量铁证,模型预测与 pin 阶梯实测误差 1-3pp + +用户质疑"trim 率为何这么低",直指表失配风险(旧表采自 sweep 形态:TRTLLM moe/ADP off/maxbs64)。**在 slo20 形态原地重采 pin{2,3,4,5} 阶梯**(u128,各 3.5-5.3 万步),得 `slo20_shape_table_v1.json`: + +| | sweep 形态旧表 | **slo20 形态新表** | +|---|---|---| +| fixed | 15.6ms | **35.3ms** | +| α(16/rank) | 1.8ms | 1.4ms | +| θ(96 token) | ~14.5ms(插值) | **10.5ms** | +| bs16 满窗步时构成 | — | 35.3+1.4+10.5=47.2ms,**不可裁部分 78%** | + +**pin 阶梯实测(SLO 口径,同点四格)**: + +| pin | SSE p50 tok/s | accept | steps | 模型预测 | 实测 | +|---|---|---|---|---|---| +| 5(满窗) | 91.6 | 3.61 | 34.7k | — | — | +| 4 | 88.6 | 3.41 | 35.9k | −4.3% | **−3.3%** | +| 3 | 83.5 | 3.06 | 38.6k | −11.3% | **−8.8%** | +| 2 | 63.6 | 2.25 | 53.4k | −32.8% | **−30.6%** | + +- **成本模型在 30% 动态范围内预测实测误差 1-3pp**——θ/accept 双侧定价都对。裁剪在此形态每一档都是净损,**planner 拒裁(trim 0.23%)是唯一正确答案**;新表比旧表更不看好裁剪(fixed 翻倍),即拒裁结论被强化而非推翻。 +- **根因对照**:poetry/agg 形态 54-56% 步时可裁 → 裁剪赚;slo20 形态**只有 22% 可裁**(ADP all-to-all + MEGAMOE_DEEPGEMM + 131k 窗稀疏注意力全是 per-step 固定成本)+ AA 高接受率 [0.98→0.89] → 裁剪两头亏。**"收益 = 可裁比例 × 接受曲线陡度"在两个形态上完成定量闭环。** +- 方法论再证:**表必须在目标形态原地重采**(同模型同硬件,仅 moe 后端/ADP/maxbs 不同,fixed 就差 2.3×)。 +- **sched x 新表已验证(预测命中)**:trim **0.0016**(新表比旧表更保守:0.0023->0.0016,仅 40/29k 请求被裁)、accept 3.59、SSE p50 **91.5** vs notrim 91.6/91.7 —— 逐位打平,零回归。置信度调度在本形态的正确行为就是「看表拒裁」,特性本身无损。 + +# 分析 + +## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) + +1. **host 未复现**:初测 +8.3%(poetry 128/rank)在四轮中缩为 [0.0, +1.9, +1.3, +2.2]%;arena 侧 [−2.4, 0.0, −1.4, −5.6]%。host 整体 ≈ 0±3%,poetry 微正、arena 微负。host 调度行为本身跨节点完全可复现(trim 22.4-22.7%,accept 1.66-1.67,四轮逐位一致)——**是收益不存在,不是实验不稳**。 +2. **device "+63.9%" 判定为假象**:device 臂生成损坏(见 Debug 日志 12:30 条),吞吐虚高由重复特殊 token + 零 EOS + 满批不排水制造。本轮 device 吞吐(+50%~+55%)全部作废。**H3 判 (c):接受判定/停机记账损坏**,旧 accept_stat 2.03 即其指纹。 +3. 待办:device 路径根因定位与修复(736879/736880 debug 节点已排队);修复后 device 收益需重测才有结论。H2/H4 挂起,等 device 修复。 +4. 联动结论:#28(SGLang 原生动态在 B300 高 bs 亏损)+ 本轮 host≈0 → 细刀 tier + 校准 + 修表后,**host 动态裁剪在 B300 agg 上没有可复现的吞吐收益**;与 SGLang blog 的 +10-20%(H200/H100)相比,B300 的 θ 平坦区更宽,裁剪节省的 token 成本换不回接受长度损失。 + +## 【08-11 09:40】主线 → laliao 侧:poetry 原文与构造脚本已入库(回应 disagg 复现请求) + +- **`docs/dspark_bench_assets/poetry_prompt.txt`**:终局判决(+5.9~+24.7%)所用 poetry prompt **逐字原文**(296 字节,memory/rivers/glass/time,≥8 场景,禁复用意象)。 +- **`docs/dspark_bench_assets/p25fix_bench.py`**:构造与打流客户端原件——poetry = 单 prompt 复制 fan-out ×N;arena = arena-hard-v0.1 全集 500 题按序循环(`prompt` 字段);**`max_tokens=768, temperature=0.7`**,warm rep=-1 + 5 计量 reps,指标 = 聚合 completion tokens / wall;asyncio 单线程,node-local 打 localhost(登录节点千线程会饿死自己的并发)。 +- **勘误(feature doc 同步修正)**:此前"自然 EOS ~760 tok"描述有误——实测均值 ~760 是 **768 上限基本打满**(poetry 几乎全部 length-capped),对表时请按 out768 记。 +- arena/throughput_1k 你侧 provenance 判断正确,无需我方再推(arena_questions.jsonl 即 arena-hard-v0.1 全集 500 题;parquet 路径已在 feature doc)。 +- 你侧 harness 若单 job 同 NVL segment 分配,MNNVL 坑理论上天然免疫(IMEX 就位)——但**请留 bring-up 判决为准**,失败签名是 ctx 侧 NIXL wait ~336s 后 FAILURE + 雪崩秒拒;R5b(主线,MNNVL=n + 600s/900s 超时)六对在跑,出数后两边可直接对表。 + +## 【08-11 09:50】混合步裁剪 A/B 判决(mixab 三臂,agg throughput_1k):机理实锤,门控回收损失 + +用户质疑"混合步不进 ragged path"→ 代码取证(planner 无 ctx gate,混合步照裁、eager 执行,成本模型 ctx 盲)→ 加一行实验门控 `TLLM_DSPARK_SKIP_MIXED_TRIM=1`(rank 本地批含 ctx 即 decline,现有 all-or-nothing 门自动组一致)→ 三臂同协议实测(各自独立节点,3 reps 中位): + +| 目标 bs/rank | notrim | sched | sched+跳混合步 | +|---|---|---|---| +| 32 | 1,858 | 1,710(**−8.0%**) | 1,840(**−1.0%**) | +| 128 | 3,798 | 3,865(≈0) | 3,876(≈0) | + +- **机理坐实**:t32 门控命中 399/736 步(54%,恰为混合步密度),trim 30.5%→5.7%,accept_len 1.38→1.54,吞吐损失 −8.0%→−1.0% ——**损失几乎全部来自混合步上的裁剪**(纯 gen 步的裁剪在本负载 ≈ 中性); +- 产品方向确认:agg 形态需要混合步自保(默认门控或成本表加 ctx 维度);disagg(纯 gen)天然免疫; +- **bs 配方对照(用户质疑成立)**:t128 的 maxbs=2×+mnt=16384 协议没把 bs 钉到 128(gen 步中位 ~89),且 notrim 吞吐比旧协议(maxbs=128/mnt=8192,4,336)低 ~12%——mnt 抬升让准入步更重,纯亏。两旋钮处方已从 feature doc 撤下。 + +## 【08-11 10:30】混合步判决同节点确认 + disagg 传输终局 + gen 饥饿物理(主线) + +- **混合步判决锁定**(同节点背靠背、双节点臂序对调):748120 sched 1721 vs skipmixed 1859(+8.0%);748121 反序 1739 vs 1855(+6.7%);skipmixed ≈ notrim(1858)打平。t32 的损失全部来自混合步裁剪,`TLLM_DSPARK_SKIP_MIXED_TRIM` 完全回收。 +- **disagg 传输六轮终局(R6 全绿,fails=0)**:最后一刀是 **`UCX_TLS=cuda_copy,cuda_ipc,sm,self,tcp`**(去 IB verbs;v2 manager 的 VMM/fabric 显存走 IB RDMA 注册有问题,MNNVL=n 必要但不充分——金丝雀"零失败"是提前杀进程的假阴性,R5b 同配置 40-80s 断开雪崩)。CI 对 GB300 的 UCX_TLS 处方与此完全一致(test_disaggregated.py)。完整配方五件套:pmix + NIXL/PYTHON + MNNVL=n + UCX_TLS 去 IB + 600s/900s 超时。 +- **gen 饥饿物理**:throughput_1k(ISL:OSL≈17:1)下 1 ctx 喂不饱 1 gen——gen 实测 bs 仅 1-5/rank,1:1 矩阵测的是 ctx 管道(lbs128 仅 ~1.1k tok/s)。QA 8k1k 用 ctx8:gen1;17:1 需要更高配比。**解法:客户端 OSL 旋钮**(`SPEEDBENCH_MAX_TOKENS`)拉长请求寿命——OSL=1024 时 gen bs 90 秒即爬到 22-25/rank 并继续上行,S/N 128 对的配对补测(同调用 3 稳态 reps)在收。 +- 副产品:OSL 旋钮同样解锁了 **agg 上实测 bs=128 的表采集**(v3 当年 128 档采不到的"早完成者"问题源于 OSL=64 寿命太短)——v4 采集(OSL 512、client 1024、pin 5/4/3/2)已在 mixab 空闲节点上跑。 + +## 【08-11 10:50】disagg 首个有效数据格(OSL1024,128 对配对稳态):纯 gen 浅 bs 仍 −8.8%;统一律成形 + +- **数字**(饱和管道 3 稳态 reps,0 fails):notrim 中位 17,294 vs sched 15,770 tok/s = **−8.8%**,三 rep 全序一致(跨对比较,各自全新 setup)。 +- **机理对齐**:bench 窗口 gen 实际 bs 仅 **29-49/rank**(1 ctx 喂 1 gen 的上限,即使 OSL=1024);这落在 B300 θ 平坦区(#28 的负转平边界 32-64/rank)。sched 裁了 12.3% token(32% 步选低档),regret 18.9%,accept 2.27(书文续写高接受,被裁 token 更贵)——**纯 gen 步不是充分条件,bs 深度才是**。 +- **统一律**:裁剪收益 = f(纯 gen 步密度 × bs 深度)。agg×本负载:混合步吃掉收益(门控可回收到 0);disagg×本负载:纯 gen 但浅 bs,照亏;poetry/arena×agg burst:纯 gen 且深 bs(128-256/rank 满批),+5.9~+24.7%。三个 regime 一张图。 +- **下一步(节点已到期,明日执行)**:disagg 要进"顺风局"需 **ctx≥2:gen1 + OSL 1024**(把 gen bs 喂到 128/rank)+ **v4 实测表**(OSL 杠杆解锁的 agg bs=128 采集在跑);那一格才是"disagg 能否翻正"的真判决。 + +## 【08-11 11:40】v4 实测表 + 混合步门控:agg t128 翻回打平/微正(当日收官) + +- **v4 表**(`align_pro_table_v4.json`,与 v3 的 32/64 样本合并拟合):OSL=512 长寿命客户端在 agg 稳住 bs=128 纯 gen 步(1,385 iter 实测样本,"早完成者"问题根治),θ 实测到 768 token 工作点。**v3 线性外推在 M=768 高估 ~17%**(≈63 外推 vs 53.8ms 实测)——bs128 过度裁剪的定量来源;α(128)=5.0ms 补齐。 +- **终局配对判决**(同节点背靠背、双向臂序,原始对齐协议 maxbs=128/mnt=8192/client=1024):notrim 4,306/4,293 vs sched(v4+`SKIP_MIXED_TRIM`)4,322/4,365 = **+0.4%/+1.7%**;对照原矩阵同格 −1.2%(v2 表、无门控)。**修正表 + 门控把本负载 agg 从负拉回打平/微正**;上限受限于 50% 混合步不可裁 + 纯 gen 步等效 bs ~89 的浅 θ。 +- t32 档未复测(门控主导,预期 ≈0);disagg 深 bs(ctx≥2:gen1 + OSL1024 + v4)仍是唯一可能的显著正收益格,待明日节点。 + +## 【08-12 00:10】主线 → laliao 侧:arena burst 语义对表 + disagg gen 提交路径 provenance(回应两问) + +**① arena 臂语义**(profinal,+6.8/+5.9% 的那轮)——原件都在 `docs/dspark_bench_assets/`: +- 客户端 `p25fix_bench.py` 逐字:arena-hard-v0.1 全集 500 题循环取 `prompt`;`max_tokens=768, temperature=0.7`,**无 ignore_eos**(自然停止允许,实测均值 ~760 ≈ 多数打满 cap);burst-drain = asyncio.gather 一次性 N 发全下、全完成计 wall;每 (dataset,bs) 首 rep 为 warm(rep=-1 丢弃)+ 5 计量 reps;指标 = Σcompletion_tokens/wall。 +- serve:`pro_sched_arm.yaml` 原件已推;graphs [1,8,32,64,128]+padding(非全阶梯),maxbs 256、mnt 4096。 +- 你 arena ≈0 vs 我 +6.8/+5.9 的候选还有一个:**臂间同节点纪律**——arena 的节点噪声你自己测过 3.8pp(08-09 工程沉淀),profinal 是同节点臂序对调;若你 arena 双臂跨节点,±4pp 即可吞掉信号。 + +**② disagg gen 提交路径 provenance**(你的 no-commit 回归协查): +- 我方 disagg gen = **sqsh `dspark-p0/trtllm.sqsh` + 全量仓库 mount**(非 overlay):R6/osl1k 轮 = 分支树 ≈ `9c5d203766`+门控补丁;今晚 deep 战役 = `1a48727867`。 +- **反证数据点:我方分支 HEAD 的 disagg gen 提交路径是通的**——08-11 上午 osl1kb 轮,gen 完整消费 ctx 交接并产出 1024-token 完成(每 rep ~1.04M token,fails=0,KV util 正常起落)。你的症状(verify+draft 真执行、active 恒 64/512、零完成、KV 平线)在我方全分支历史中未出现过。 +- 我方近期 commit 无有意改动 disagg gen 首步提交/接受路径(门控补丁只在 planner phase-1 加 decline)。你的 base 是 `faf2c60935`+28 文件 overlay——**怀疑点排序:overlay 组合(部分文件新、部分旧的 Franken 树)> sqsh 基底 > 分支本体**;D1 判别(旧镜像无 overlay)正是对的刀口。若 D1 通,建议直接换"全量 mount 我们分支 HEAD"复跑,跳过 28 文件二分。 +- 我方 deep 战役(3ctx:1gen,OSL1024,今晚)同路径在跑,出数即再添一个分支 HEAD 的通/不通数据点。 + +## 【08-12 02:00】深 bs disagg 终局(throughput_1k,3ctx:1gen,OSL1024,v4 表):统一律最后一格点亮,+4.4%(同节点配对) + +- **设计**:3 ctx(tp8)喂 1 gen(DEP8),OSL=1024 拉长寿命(稳态 gen bs 中位 116/rank,120-128 区间大量出现,吞吐为 1:1 disagg 的 10 倍),client 2048(2× 超发),两 quad 臂序对调,五件套 env,v4 实测表。 +- **数字**(稳态 3 reps,0 fails): + - 同节点配对(gen 749968,N→S):notrim med 37,588 vs sched med 39,247 = **+4.4%**; + - 跨 quad:quad2 sched med 43,242(其 notrim 臂死于 walltime)——表面 +15.0%,但同臂跨 quad 差 ~10%(43,242 vs 39,247)暴露强节点效应,**该数字不可引用,以同节点 +4.4% 为准**。 +- **判读**:throughput_1k 在"纯 gen 密度 × bs 深度"双高角落**首次转正**——统一律四格闭合(同负载:agg 混合步 −8%→门控 0%,disagg 浅 bs −8.8%,disagg 深 bs **+4.4%**)。幅度低于 poetry(+24.7%)符合预期:本负载 accept 更高(书文续写 2.3-2.6,被裁 token 更贵)且 sched 3 reps 内有爬升(36,965→40,773),更多 reps 可能收敛到更高中位。 +- 方法论注记:sched 臂 rep 间上行趋势 + 双臂 range 有重叠,+4.4% 的置信度为"方向可靠、幅度 ±3pp";加固需 n≥5 双向臂序(节点已到期,列为可选补测)。 + +## 【08-12 06:00】256 档深 bs disagg 终局:+7.0%/+15.6%(双向同节点配对);v5 全实测表;bs 深度趋势完整 + +- **v5 表交付**(`align_pro_table_v5.json`,disagg 256 满批 pin 阶梯 + v4 样本合并):θ 实测覆盖 96→1536 token(θ(1536)=106.4ms),α(256)=14.4ms——32/64/128/256 全档无外推。 +- **数字**(4ctx:1gen,OSL1024,client 2048=1×,v5 表,0 fails,同 gen 节点背靠背、两组臂序对调): + - quint2(N→S):notrim 38,561 vs sched **44,561 = +15.6%**; + - quint1(S→N):sched 45,551 vs notrim 42,574 = **+7.0%**; + - sched 跨节点仅差 2%(45.6k/44.6k),notrim 波动较大(38.6k/42.6k,rep 内有爬升)→ 判读 **+7~+16%,中心 ~+11%**。 + - 实测 gen bs:notrim med 194(p90 236)、sched med 174——sched 排水快在飞更低,收益是在更浅 bs 下拿到的(锁同 bs 只会更高)。 +- **bs 深度趋势(throughput_1k,同代码同负载)**:29-49/rank **−8.8%** → ~116/rank **+4.4%** → ~174-236/rank **+7~+15.6%**——收益随深度单调走强,统一律的 bs 因子完整成像;256 档已接近 poetry 满批量级(+24.7%)。 +- **新坑入册(backlog)**:client 在飞 4096 > router(uvicorn)默认 backlog 2048 → 半数连接挂死、管道爬行(91% 失败,gen 只 40/rank)。128 档 client 恰好 2048 纯属侥幸。修法:client ≤2048,或给 `trtllm-serve disaggregated` 加 backlog 旋钮(上游项)。被污染的 4096 行已在 summary 以 tag 区分(*b = 干净复测)。 diff --git a/docs/source/features/dspark-confidence-scheduling.md b/docs/source/features/dspark-confidence-scheduling.md new file mode 100644 index 000000000000..d33e5012954c --- /dev/null +++ b/docs/source/features/dspark-confidence-scheduling.md @@ -0,0 +1,611 @@ +# DSpark 置信度调度验证（Confidence-Scheduled Verification） + +DSpark 每个 decode step 都会 draft 出一整个 block，并把它全部送给 target 验证。 +它的 draft 网络本来就带了一个 **confidence head**，为每个 draft 位置打分，预测该位置 +能否活过验证。置信度调度把这个 head 接进了验证调度器：block 依然完整 draft，但**真正 +送去验证的位置数**由这些分数决定。 + +本文讲设计。用户可配置项见[配置](#配置)；DSpark 在各投机算法中的位置见 +[Speculative Decoding](./speculative-decoding.md)。 + +## 目录 + +- [问题与核心洞察](#问题与核心洞察) +- [三条设计不变量](#三条设计不变量) +- [模块分层](#模块分层) +- [一个 decode step 的数据流](#一个-decode-step-的数据流) +- [relay 滞后](#relay-滞后) +- [算法](#算法) +- [一个完整的数值例子](#一个完整的数值例子) +- [Ragged 验证](#ragged-验证) +- [CUDA graph 处理](#cuda-graph-处理) +- [跨 rank 一致性](#跨-rank-一致性) +- [退化路径](#退化路径) +- [配置](#配置) +- [性能天花板与已知局限](#性能天花板与已知局限) +- [部署要求：DSv4 必须显式设 max_seq_len](#部署要求dsv4-必须显式设-max_seq_len否则并发被钉在-maxbs2) +- [基准配置（可复现）](#基准配置可复现) + +## 问题与核心洞察 + +基线行为： + +```text +draft: 一次并行 backbone forward -> [t1 t2 t3 t4 t5 t6 t7] (K = 7，固定成本) +target: 验证全部 7 个 + 1 个 bonus = 8 个 token 位置 +``` + +如果某个请求的 draft 在位置 3 就已经断了，那么花在位置 4-7 上的验证算力全是浪费。 +confidence head 预测的是每个位置的**条件**接受概率 `P(accept_k | accept_1..k-1)`， +因此前缀存活概率就是累积乘积： + +```text +survival[r][j] = prod_{i <= j} conf[r][i] # 位置 j 被"走到 且 被接受"的概率 +``` + +由于 `conf` 是 sigmoid 输出、落在 `(0, 1)` 内，**survival 沿位置严格非增**。 +这条单调性是后面整个分配算法的地基。 + +## 三条设计不变量 + +下面所有取舍都可以从这三条推出来： + +| # | 不变量 | 推论 | +|---|--------|------| +| **I1** | block 永远完整 draft | draft 是一次并行 forward，砍长度省不了钱。只有送给 target 的 token 数可裁 | +| **I2** | 接受规则完全不动 | 调度只决定**验多少**，绝不决定**接不接受**，因此**输出分布不变** | +| **I3** | 决策在 host 上做，且必须落在已捕获的 graph 上 | CUDA graph key 由 host 值构成；未捕获的 draft 长度不会报错，而是静默掉进 eager 执行 | + +**I2 是最重要的性质**：这不是精度/吞吐的 tradeoff，而是纯粹的算力回收。所以这个特性的 +精度测试断言的是与 baseline **完全相同**的分数 —— 任何偏差都是不变量被破坏，而不是调参结果。 + +## 模块分层 + +```text +配置层 + llmapi/llm_args.py :: DSparkDecodingConfig + +DEVICE 侧（在 target 的 captured CUDA graph 内） + _torch/models/dspark/heads.py :: DSparkConfidenceHead + forward() [*, K, hidden] -> [*, K] 原始 logits（fp32 matmul） + apply_sts() 原始 logits -> 标定后的概率（逐位置温度） + _torch/speculative/dspark.py :: DSparkWorker._draft_gen_block_batched + _confidence_logits[slots] = conf # slot 索引，原地写 + | + | 非阻塞 D2H，滞后（见下文） + v +HOST 侧（决策，完全在 capture 区之外） + _torch/speculative/dspark_verify.py :: DSparkVerifyPlanner # 协调器、退化守门 + | + +-- dspark_schedule.py survival（cumprod）+ 全局 top-k 分配 + +-- dspark_planner.py SpsCostTable + Theta = tau / T 的 argmax + tier 推导 + _torch/pyexecutor/py_executor.py :: _dspark_confidence_draft_len # 每步 hook + +批次成形层 + _torch/pyexecutor/model_engine.py ragged 输入布局、graph 形状拟合 + _torch/pyexecutor/cuda_graph_runner.py graph key 上的 token 轴 + _torch/speculative/dspark_ragged.py qo_indptr 打包、bucket 对齐、scatter + _torch/attention_backend/sparse/dsa.py indexer 元数据按 token 展开 +``` + +## 一个 decode step 的数据流 + +```text + HOST | DEVICE（captured graph） +-------------------------------------------------------------------------------- + step t | + _dspark_confidence_draft_len() | + 1. 读取滞后的快照 | + （cudaEvent.query()，非阻塞） | + 2. 按 slot 查出每个请求对应的行 | + 3. apply_sts -> cumprod -> survival | + 4. argmax Theta -> tiers 中的 draft_len | + 5. 跨 rank allgather(max) | + | + _prepare_tp_inputs() | + 铺 input_ids / position_ids | + _attach_ragged_verify_layout() | + | + graph_key = (bs, draft_len, ..., [bucket]) | + ------------- replay --------------------> | target forward + | sample_and_accept_draft_tokens + | draft forward（完整 block K） + | confidence head + | _confidence_logits[slots] = conf (*) + 6. stage_confidence(buffer, generations) | + 轮转 relay ring，然后发起 | + 非阻塞 D2H <---------------+ +-------------------------------------------------------------------------------- +``` + +## relay 滞后 + +调度决策跑在一个**滞后**的快照上，而且这个滞后是**钉死的常量**，不是"引擎给多少算多少"。 +DSpark 规定了两步回溯（two-step-back）的 relay；SGLang 也把这个值钉死，并在其 overlap +scheduler 会给出不同值时用 ring buffer 补齐。 + +钉死很重要，因为 STS 标定表和 SPS 成本表是在**某一个特定滞后**下拟合出来的。如果有效滞后 +会随一个无关的引擎开关漂移，那么一张 profile 好的表就会静默失效。 + +TensorRT LLM 这两步是怎么来的 —— 调度 hook 跑在 forward **之前**： + +```text +iteration t: [hook(t)] -> [prepare] -> [forward(t)] + | 读取 hook(t-1) staged 的快照 + | 发起 D2H，流序在 forward(t-1) 之后， + | 所以抓到的是 forward(t-1) 的 confidence + v + hook(t) 读到的是 forward(t-2) 的 confidence ==> 引擎滞后 = 2 +``` + +无论 overlap scheduler 开不开都成立：overlap 改变的是拷贝**何时落地**，而不是它抓的是 +哪一次 forward。没落地的拷贝会退化成"验全块"，而不是退化成一个陈旧的猜测。 + +实现中滞后不可配置：executor loop 恒定给出两步（上图），planner 直接消费该快照。快照行的新鲜度由 `_confidence_stamp` 守卫 —— 逐 buffer 行的 draft-pass 序号，随 confidence 经同一个 `slots` scatter 在 graph 内写入；陈旧或从未打分的行读到中性 confidence，survival 归 1.0，退化为验证整块。 + +### slot 索引与 generation 守卫 + +快照按 **worker slot** 索引，绝不按 batch 位置。请求会在"写入分数的那一步"和"读取分数的 +那一步"之间进出，所以位置 `i` 经常已经是另一个请求了。 + +slot 还会被回收，这带来两种不同的陈旧风险，由两个互补机制处理： + +| 风险 | 机制 | +|------|------| +| device 上那一行还留着上一个占用者的分数 | `_assign_slot` 在把行交给新请求时重置为中性值 | +| host 侧快照里还留着回收之前的副本，device 侧的重置够不着 | 每行带一个 `_confidence_stamp`（draft-pass 序号），与 confidence 一起 stage；与当前 pass 对不上就把该行按中性处理，survival 归 1.0 | + +这两者是**互补的，不是二选一**。generation 标签是在 stage 时于 host 侧快照的，而 +confidence 是在 forward 时于 device 上写的 —— 所以正是那次重置让标签说的是真话：只 bump +不重置，就会把上一个占用者的分数打上新占用者的标签。反过来，只重置也清不掉快照里已经 +存在的 host 副本。 + +两种风险都收敛到同一个安全方向：**survival = 1.0，即验证整个 block**，也就是启用该特性 +之前的行为。 + +## 算法 + +### 第 1 步 —— survival + +```python +survival = torch.cumprod(sigmoid(logits / sts_temperatures), dim=1) +``` + +标定（STS）之所以重要，是因为调度器消费的是**累积乘积**：逐位置的偏差会沿 block +**几何级放大**，未标定的 head 会系统性地高估前缀存活率。 + +### 第 2 步 —— 成本模型 + +本实现遵循已公开的形式： + +```text +T(bs, K) = bias + alpha(bs) + theta(M) + ---- --------- -------- + 固定 批次相关 verify-token 成本 + （draft pass、 ^ 唯一可裁剪的一项 + 权重搬运） +``` + +`SpsCostTable` 按总验证 token 数 `M` 做**钳位线性插值**查表（对齐 SGLang 的加法表消费端）。想表达"平台"就必须把平台测出来——两个等值断点；消费端不再替硬件假设断点之间是平的（那个假设曾把 1512 token 的满块按 768 的价格计费，整个特性因此从未裁剪过）： + +```text +step_time + | +-------- <- riser：新的 kernel wave 起来了 + 9 | | + | +-----+ + 4 | | <- shelf：这一段上加 token 是真的免费 + 2 |--+ + +--+-----+-----+----> M（总 token 数） + 8 16 32 +``` + +**token 口径**：一个请求验证 `L` 个 draft 位置，实际提交 `L + 1` 个 token —— 它手上那个 +bonus token 加上它的 draft。所以 `M = bs * (L + 1)`。这个换算只在一个地方计算 +（`dspark_planner.total_verify_tokens`），保证 planner 和 tier 推导不会漂移。 + +由于 `Theta = tau / T` 是**比值**，不可裁剪的那几项**不会约掉**。低估它们会让每个验证 +token 显得比实际更贵，planner 就会过度裁剪。因此这些项是**给进来的，不是猜的**： +`fixed_overhead_ms` 承载 `bias`，可选的 `batch_sizes` / `batch_overhead_ms` 阶梯承载 +`alpha(bs)`。 + +### 第 3 步 —— 目标函数与全局 argmax + +```text +tau(L) = num_gen_requests + sum_r sum_{j L +``` + +### 第 4 步 —— shelf 右边缘性质 + +> 在一个 cost shelf 内，step time 恒定，而 `tau` 随每个新增 token 严格增， +> 所以 `Theta = tau / cost` 在 shelf 上严格递增。因此最优点**永远落在 shelf 的右边缘**， +> 绝不在其内部。 + +于是由各 shelf 右边缘构成的 tier 集合，**在推导时所用的那个 batch size 下**精确包含连续最优解： + +```text +right_edge(breakpoint T) = (T - 1) // bs - 1 # length 空间 +``` + +这个**零损失性质跨 batch size 不成立** —— cost shelf 活在 token 空间，所以 shelf 右边缘 +在 length 空间是 `bs` 的函数。一个部署会捕获很多 batch size 但只能捕获一套 ladder，因此 +应当在**稳态众数 batch size** 上推导，并**实测**其他 batch size 上的残差损失，而不是假装它不存在。 + +### 第 5 步 —— 逐请求分配（ragged 模式） + +预算是 `num_gen * (L* - min_verify_len)` —— 和 uniform 决策**花掉的 token 总量相同**， +只是重新分配。 + +`min_verify_len` 在本实现中**固定为 1**（配置校验强制 >= 1）：每个请求每步至少验证 +anchor 位置。SGLang 允许 floor 为 0——draft 已死的请求拿 0 个位置也不是饿死（bonus token +不在 confidence 矩阵里，该请求照样出 1 个 token，退化为普通自回归），省下的预算可以喂给 +强请求的深位置，混合批下其可达 Θ 严格更优（test_starving_weak_rows_is_a_known_gap 钉死了 +这个差异）。移植 floor=0 需要 executor 支持零窗口请求（layout/采样路径改动），是已知的 +待决 gap 而非本实现的默认。 + +```python +candidates = survival[:, floor:].flatten() # 每一个 (request, position) 对 +# 按 survival 全局排序，取前 budget 个，然后按行计数 +``` + +不需要任何显式的前缀约束：survival 沿位置非增，所以**任何全局 top-k 自动就是每个请求的 +前缀**。因此分配逻辑只需要**数**每行中了几个，不需要追踪中的是哪几个。 + +tie 用 `(position, request)` 打破，**绝不用数值**。两个 TP rank 即使算出 bitwise 不同的 +confidence，也必须选出相同的 verify 长度，否则它们的 batch 形状（进而 collective）会发散。 + +## 一个完整的数值例子 + +`K=7, bs=4, tiers=[1,3,7], min_verify_len=0`（注意：这是 SGLang 语义的演示例；本实现 floor 固定为 1），成本表 +`token_counts=(0, 8, 16, 32)`、`step_time_ms=(2.0, 2.2, 2.6, 9.0)`。 + +survival 矩阵： + +```text + j=0 j=1 j=2 j=3 j=4 j=5 j=6 + r0 | 0.95 0.87 0.79 0.69 0.59 0.47 0.35 <- draft 很健康 + r1 | 0.60 0.18 0.02 0.00 0.00 0.00 0.00 <- 第 2 位就死了 + r2 | 0.90 0.81 0.73 0.66 0.59 0.53 0.48 <- 健康 + r3 | 0.70 0.49 0.34 0.24 0.17 0.12 0.08 <- 中等 + 列和 3.15 2.35 1.88 1.59 1.35 1.12 0.91 +``` + +uniform 决策： + +| L | M = 4*(L+1) | T (ms) | tau = 4 + 列和 | Theta | +|---|-------------|--------|----------------|-------| +| 1 | 8 | 2.2 | 7.15 | 3.25 | +| **3** | **16** | **2.6** | **11.38** | **4.38 <- 最大** | +| 7 | 32 | 9.0 | 16.35 | 1.82 | + +于是 `runtime_draft_len = 3`，每个请求验证 3 个 draft 位置。 + +ragged 用**同样的预算**（`4 * 3 = 12` 个位置）重新分配。floor 固定为 1（`min_verify_len >= 1`，每个请求至少验证 anchor 位置；SGLang 允许 floor 为 0、可把绝望请求饿到零窗口，这是两者的已知语义差异），floor 之上的位置参与全局竞争，排序后取预算内的前若干个： + +```text +0.95(r0) 0.90(r2) 0.87(r0) 0.81(r2) 0.79(r0) 0.73(r2) +0.70(r3) 0.69(r0) 0.66(r2) 0.60(r1) 0.59(r0) 0.59(r2) + \--------------- r0 拿 5 个，r2 拿 5 个，r3 和 r1 各 1 个 ---------------/ +r3 的下一个候选 0.49、r1 的 0.18 都落在预算之外 +``` + +| | verify_len（draft 位置数） | 期望产出 | +|---|---|---| +| r0 | **5** | 3.89 | +| r1 | **1** | 0.60 | +| r2 | **5** | 3.69 | +| r3 | **1** | 0.70 | +| **合计** | **12**（与 uniform 的 `4 * 3` 完全相同） | **tau = 12.88** | + +```text +uniform L=3 : 12 个验证位置 -> tau = 11.38 +ragged : 12 个验证位置 -> tau = 12.88 （成本完全相同，+13%） +``` + +这就是 ragged 模式的全部动机：把 r1 和 r3 注定要浪费掉的位置，转移给 r0 和 r2。 + +## Ragged 验证 + +以 `verify_lens`（token 单位，含 bonus）`= [6, 2, 6, 2]` 为例： + +```text +qo_indptr = [0, 6, 8, 14, 16] # 独占前缀和 + +flat token 轴（attention kernel 想要的布局）: + idx 0 1 2 3 4 5 | 6 7 | 8 9 10 11 12 13 | 14 15 + b d0 d1 d2 d3 d4 | b d0 | b d0 d1 d2 d3 d4 | b d0 + \------- r0 ---------/ \- r1 -/ \-------- r2 ---------/ \- r3 -/ + ^ ^ ^ ^ + qo_indptr[0]=0 [1]=6 [2]=8 [3]=14 +``` + +### bucket 对齐 + +被捕获的 token bucket 是每个 tier 对应的 `bs * (t + 1)`。一个 batch 把自己的总数向上取整 +到其中之一，`RaggedVerifyLayout.fill_bucket` 分两阶段花掉这部分余量：**先给真实请求** +（上限 `max_verify_len`，这样这些 token 去验证的是 step 本来就要付钱的真实 draft 位置）， +剩下的才给 pad 行。 + +token 数必须**精确**命中 bucket。`seq_lens.sum()` 会变成 `attn_metadata.num_tokens`， +它会在 attention-DP 各 rank 间 all-gather，并驱动 MoE 的 chunk 数。如果少报，attention 和 +MoE 对"在飞多少 token"的认知就不一致 —— **而且什么都不会报错**。 + +### 形状恒定，内容可变 + +这正是 ragged batch 能被捕获的原因： + +```text +graph 捕获的是 -> (padded_bs 行, bucket 个 token) # 一小组固定形状 +每步变化的是 -> verify_lens / qo_indptr 里的数值 # 持久 buffer +``` + +## CUDA graph 处理 + +```text +uniform: (bs, draft_len, False, short_seq_mode, is_all_greedy) + ^ tiers 中的某一个 + +ragged: (bs, top_tier, False, short_seq_mode, is_all_greedy, bucket) + ^ token 轴 +``` + +ragged 下 `draft_len` 钉死在最高 tier（不变量 I1：block 总是完整 draft），变化的是 token +总数。两种模式都捕获 `|batch_sizes| * |tiers|` 张图。 + +与"按 batch size 调度 draft 长度"不同，这里的 tier **不是** batch size 的函数 —— 在固定 +`bs` 下，随着 batch 的 confidence 变化，调度器仍可能选不同的长度 —— 所以捕获集合是**叉积**。 +每张捕获的图大约花 10-23 MB 元数据，而这块内存**直接从 KV cache 里出**（KV pool 是按 capture +之后剩下的量来定的）。所以 ladder 必须短。 + +capture 安全性依赖几个具体选择： + +| 选择 | 原因 | +|------|------| +| `return_confidence` 在构造时读一次，绝不逐步判断 | 否则捕获的 draft graph 会发散 | +| STS buffer 原地更新，绝不 rebind | `nn.Module.__setattr__` 换 buffer 会换 `data_ptr`；先前捕获的 graph 会一直读旧存储，标定静默失效 | +| `_confidence_logits` 是持久 buffer，通过 slot 散射写入 | replay 时落到当前 batch 实际占用的槽位 | +| ragged buffer 通过持久分配的切片写入 | 保持 `data_ptr` 稳定，graph 才看得见新值 | +| `repeat_interleave` 一律传 `output_size` | 否则 torch 要把 repeats 的累积和读回 host 来确定输出尺寸 —— 在 capture 期非法，在其他地方则是每步一次同步 | + +## 跨 rank 一致性 + +```text +draft_len 是 CUDA graph key 的一部分 + | + v +但它（当时）不在 attention-DP 的一致性 allgather 里——这是促成现行 shape gate 的 +历史事故：现在的实现已把 draft_len 纳入 allgather 并逐 rank 比对，下图保留为动机说明 + | + v +两个 rank 选了不同长度 -> 选到不同的 graph + | + v +一个 replay、一个掉进 eager -> 它们的 collective 发散 +``` + +归约取的是 **max 而不是 min**：想裁更多的 rank 只是多验几个 token，而 graph key 不一致 +是**不可恢复**的。ragged 下捕获形状是二维的，一致性来自"对归约结果取整"，而不是"各 rank +对自己的数取整"： + +```text +padded_bs = round_up_bs( max_r num_real_r ) +bucket = round_up_bucket( padded_bs + max_r slack_r ) +``` + +## 退化路径 + +静默退化是这个特性的主要失效模式，所以每一条放弃裁剪的路径都在 `planner.stats` 里计数： + +```text + +- planner 未构建 -----------------> 全块（无 checkpoint） + | + +- 成本表 flat（仅当只给了 sts 表时可达；两表 + | 皆缺在构造时即被拒）------------> 全块 + 告警 + | "每个 token 都免费" => 裁剪等于盲猜 + | +_dspark_confidence --+- 快照还没落地 -------------------> 全块 +_draft_len() | （cudaEvent.query() == False） + | + +- 快照行数少于 batch --------------> 全块 + | + +- 请求从未被打分 ------------------> 中性行 -> 全块 + | （刚 prefill 完，或 slot 刚回收） + | + +- 找不到已捕获的 ragged 形状 ------> 退回 uniform（uniform 一定有图） +``` + +所有路径都朝同一个方向退化：**验证整个 block**，即启用该特性之前的行为。输入不可信时， +调度器**永远不会**比静态基线验得更少。 + +部分结果被当作失败而非"部分成功"：如果 planner 无法给出恰好每请求一个 verify 长度，整批 +退回 uniform。一个"半窗口"的 batch 比完全不加窗口更糟 —— 输入布局是逐请求构建的（会走成 +ragged），而 spec metadata 看到缺失的窗口后会保持 uniform。 + +## 配置 + +```yaml +speculative_config: + decoding_type: DSpark + max_draft_len: 7 + + # 主开关。block 仍然完整 draft，只改变送给 target 验证的 token 数。 + enable_confidence_scheduling: true + + # 逐位置 STS 标定。省略 => 裸 sigmoid（恒等标定）。 + # 调度器消费的是累积乘积，所以逐位置误差会沿 block 几何级放大。 + confidence_sts_path: /path/to/sts.json + + # profile 好的 step-cost 曲线。开启 ragged 时必填（除非提供 confidence_sts_path）： + # 两者皆缺在构造时即报错 —— flat 成本模型下 planner 永远拒绝裁剪，跑完也是空跑。 + confidence_sps_table_path: /path/to/sps.json + + # 被捕获的长度 ladder；调度器只能从这个集合里选。 + # 默认：[1, ceil(K/2), K] + confidence_verify_len_tiers: [1, 3, 7] + + # 逐请求 verify 窗口。必须与 enable_confidence_scheduling 同开 —— 两个 flag 描述同一个 + # 特性，只开其中一个会在构造时被拒绝；关闭特性 = 两个都设 false。 + enable_ragged_verify: true +``` + +`sps.json`： + +```json +{ + "token_counts": [0, 128, 256, 512], + "step_time_ms": [2.0, 2.1, 4.0, 9.0], + "fixed_overhead_ms": 0.0, + "batch_sizes": [1, 16, 64, 256], + "batch_overhead_ms": [3.0, 5.0, 17.0, 40.0] +} +``` + +`token_counts` 按**总验证 token 数**索引，**包含 bonus token**。`fixed_overhead_ms` 和 +`batch_*` 这一对是可选的，但绝不是装饰 —— 见[算法](#算法)中的成本模型部分。 + +## 性能天花板与已知局限 + +```text +一个 decode step = bias + alpha(bs) + theta(M) + ------------------ -------- + 调度动不了这部分 唯一可回收的一项 +``` + +因此收益是一个**高并发效应**。在小 batch 下，step 成本对验证 token 数几乎是平的，裁剪省 +不下时间，却仍要付出接受长度的代价 —— 打平甚至为负。 + +**在目标硬件上量出可裁剪部分的占比是前提，不是事后工作。** 如果 `alpha(bs)` 压过 +`theta(M)`，那么无论 confidence head 多好，天花板都很低。两个测量陷阱： + +- 在**固定 batch size** 下 profile，会让 `M = bs * (L + 1)` 与 `L` 完全共线，`alpha` 和 + `theta` 分不开。必须扫 batch size。 +- 如果 `d(step)/dK` 在各 batch size 上是平的，说明成本是 per-draft-pass 而非 + per-verify-token —— 调度器几乎没有可回收的东西。 + +其他局限： + +| 局限 | 说明 | +|------|------| +| tier ladder 跨 batch size 有损 | shelf 右边缘随 `bs` 移动，一套 ladder 不可能同时坐在所有 batch size 的边缘上。残差要实测 | +| 快照是 2 步之前的 | 纯吞吐启发式（I2 保证正确性与它无关），但在剧烈抖动的负载下决策质量会下降 | +| ragged 模式的端到端数字仅在一个平台上取得 | DeepSeek-V4-Pro-DSpark，8×B300 DEP8，overlap + CUDA graph（GSM8K 96.44，基线 96.21）；仍默认关闭，跨硬件需重测 | +| 逐 token 采样参数路径只有非全贪心 batch 才会触发 | `temperature=0` 的精度跑覆盖不到它 | + +## 部署要求：DSv4 必须显式设 `max_seq_len`（否则并发被钉在 ~maxbs/2） + +DSv4 的窗口池组（SWA + compressor KV/score，窗口 = 128 + `max_draft_len`，DL=5 时 133 +token）在 v2 KV manager 里是固定槽分配器。不显式设 `max_seq_len` 时，serve 会从模型推断 +出 1M 并把它当 `typical_seq_len` 喂给比例定容——全注意力池吃掉全部配额，窗口池被饿到下限 +`min_slots ≈ 1.05 × max_batch_size` 块（`kv_cache_manager_v2.py:1775-1779`），而下限按 +history=0 建模成每请求 1 块（`:1803-1812`）。运行时窗口 133 > tokens_per_block 128，每个 +生成请求实际锁 2（偶 3）块（`_life_cycle_registry.py:38-51`），于是可调度并发 += 1.05·maxbs / 2.07 ≈ **0.51 × max_batch_size**（实测：maxbs=128 → 61-63 行；maxbs=256 +→ 131-136 行），伴随 ~2.7× 的步时惩罚。关闭 spec 时同一机制不降级而是**调度器死锁** +（`scheduler_v2.py:439` 处崩溃，报错信息有误导性）。 + +**修法：serving 配置里显式写 `max_seq_len`（按真实负载，比如 8k1k 场景写 9216）。** +`typical_seq_len` 随之落地，比例定容给窗口池的份额远高于下限，满 maxbs 可调度。判别实验 +证明只动这一项即解钉；只调 KV fraction / max_num_tokens 均无效。历史工作绕法 +`max_batch_size = 2 × 目标并发` 是同一算术的另一侧（把下限抬高一倍），在显式 +`max_seq_len` 之后不再必要，但对 ctx 准入 headroom 仍有独立价值（见下）。 + +代码侧的根治（`typical_seq_len` 不应默认推断的 max_seq_len、下限按稳态 2 块建模、 +no-spec 应降级而非死锁、字节估算器补 draft 窗口放大）待报 upstream，不阻塞部署规约。 + +## 基准配置（可复现） + +完整实验记录与逐轮数据见 `docs/dspark_sts_program.md`。以下为两类负载的精确配置。 + +### 负载 A：poetry / arena（agg，burst-drain；收益 +5.9~+24.7% 的来源） + +- 硬件/形态：DeepSeek-V4-Pro-DSpark，8×B300 单节点，`--tp_size 8 --ep_size 8`（DEP8， + `enable_attention_dp: true`），`--max_batch_size 256 --max_num_tokens 4096`。 +- 数据集： + - **poetry**：单条固定英文 free-verse 长诗 prompt（逐字原文见 + `docs/dspark_bench_assets/poetry_prompt.txt`），复制 N 份。短输入（~60 tok）+ + `max_tokens=768, temperature=0.7`（实测均值 ~760 tok/req，基本打满上限）—— + 短 ISL / 长 OSL 的极端。 + - **arena**：arena-hard-v0.1 全集 500 条真实问题（jsonl，取 `prompt` 字段，N>500 时 + 循环复用），同 `max_tokens=768, temperature=0.7`。中等 ISL / 长输出的多样负载。 + - 打流客户端原件：`docs/dspark_bench_assets/p25fix_bench.py`。 +- 客户端 burst-drain：一次性下发全部 N（512/1024）个请求，计时到全部完成，n=10 轮取中位。 + **注意：该协议下 99% 的步是纯 gen 步**（prefill 波在头几十步内结束）——这是收益成立的 + regime 前提。 +- S 臂 spec 配置（`pro_sched_arm.yaml`）：tiers `[2,3,4,5]`， + `confidence_sps_table_path: postfix_pro_table.json`（修复 #31/#32 后在 maxbs=256 满批 + 重采：bs128 实测 θ(L2/3/4/5)=80.8/90.3/102.7/110.3ms），**未挂 STS**（裸 sigmoid）。 + N 臂同配置去掉两个 confidence 开关。kv 0.5，block reuse 关，graphs [1,8,32,64,128] 带 + padding，overlap on。 + +### 负载 B：throughput_1k（roofline 对齐负载，agg 持续饱和） + +- 数据集：roofline 报告使用的 **throughput_1k** parquet + （`iputterman/speed/throughput_1k/test-00000-of-00001.parquet`），按文件序取前 1024 条、 + 每条取 `turns[0]` 为 prompt（实测 ISL≈1110 tok）；greedy（temp=0），`max_tokens=64`。 + 注：源 parquet 的 SHA 与 roofline 报告冻结 dump 不同，对比其报告数字时需留意版本。 +- 客户端持续饱和：固定在飞数，完成即补位，ramp 256 个完成后计量 1024 个完成。 +- 关键物理：OSL=64 + 长 prompt + 持续补位 → **~50% 的步是 ctx+gen 混合步**，纯 gen 采的 + 成本表在混合步高估节省（成本模型 ctx 盲），裁剪收益不成立（sched −1.2~−6.3%）。收益是 + **负载 regime 的函数**，不是普适常数。 +- **bs 定量**：客户端在飞数 = 目标并发（8×目标/节点），maxbs/max_num_tokens 与生产配置 + 一致即可。本负载 OSL=64、请求周转极快，实测 gen bs 会因准入/周转间隙低于标签值——报告 + 时以 iter log 的实测中位为准，不以档位标签为准。 +- 成本表与 STS 均须在本负载重采（表按 `--from-iter-log` 纯 gen 段拟合）。 + +### 负载 B'：disagg（gen-only 吞吐；裁剪的顺风 regime） + +ctx/gen 分离后 gen server 每步纯 gen，成本表全程在域内。两台 B300 独立 slurm 作业, +NIXL 传输。**bs 钉法**：gen server 没有 prefill 争抢槽位，直接 +`max_batch_size = 目标并发` + 客户端适度超发（2–4×，够填补周转间隙即可）把 gen bs 钉在 +maxbs 上；**必须从 gen 侧 iter log 监控实测调度 bs 确认钉住**，不能只信标签。超发过猛没 +有额外收益，只会加深 ctx 队列、压穿 KV-transfer / router 超时。**四个必踩的坑与解法**： + +1. v2 KV manager 只支持 `transceiver_runtime: PYTHON` + `backend: NIXL`（C++/DEFAULT + 传输在启动时 `CUDA_ERROR_INVALID_CONTEXT`）。 +2. **必须 `export UCX_CUDA_IPC_ENABLE_MNNVL=n`（两侧）**：v2 用 + `CU_MEM_HANDLE_TYPE_FABRIC` 分配 KV,无 NVLink 互联的独立作业间 UCX 会误判 cuda_ipc + 可达 → RDMA 卡死 ~336s → NIXL 元数据雪崩（全请求失败）。`ucx_perftest` 用普通 + cudaMalloc,对此 bug 是假阴性。NVL72 机架内（IMEX 就位）不需要。 +3. 超时三层都要抬：`cache_transceiver_config.kv_transfer_timeout_ms: 600000`（默认 60s, + 超发排队下成批误杀）;router `trtllm-serve disaggregated --request_timeout 900`（默认 + 180s）;（传输层由 2 解决）。 +4. **ctx 与 gen 必须共用同一个 `speculative_config`**（YAML anchor）——否则窗口池几何不 + 一致（128/8 vs 133/13）且 draft KV 无人 prefill。 + +```yaml +# ctx 与 gen 共同的骨架（两侧一致的部分） +max_seq_len: 8192 # 见"部署要求" +speculative_config: &spec + decoding_type: DSpark + max_draft_len: 5 + speculative_model: +kv_cache_config: {enable_block_reuse: false, free_gpu_memory_fraction: 0.5} # gen 侧 0.7 +enable_attention_dp: true +cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON + max_tokens_in_buffer: 8192 + kv_transfer_timeout_ms: 600000 +# ctx 侧另加 disable_overlap_scheduler: true;gen 侧在 *spec 上打开 confidence 开关。 +``` + +## 参考 + +- DeepSeek DeepSpec： +- SGLang 的 DSpark 集成： diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py index 6dd402369d9f..bfd17faf2f28 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py @@ -221,7 +221,12 @@ def forward( if num_generations > 0: gen_kv_lens = kv_lens[num_contexts:] + # Uniform batches: every request appended next_n tokens. Ragged + # batches: next_n is only the batch maximum (the kernel's loop + # bound) and the real per-request counts come through + # new_tokens_per_seq. next_n = metadata.num_gen_tokens_per_seq + new_tokens_per_seq = metadata.gen_new_tokens_per_seq # Pass full kv_score (not sliced) with the generation portion of # cu_seq_lens so the kernel reads at the correct absolute offsets. torch.ops.trtllm.compressor_paged_kv_compress( @@ -240,6 +245,7 @@ def forward( self.head_dim, self.compress_ratio, next_n, + new_tokens_per_seq, ) # Scatter to cache with appropriate quantization (all modes fused) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index 1f76e5a73e51..65350c249485 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -334,6 +334,19 @@ def __post_init__(self): ) self.cu_seq_lens[0] = 0 + # Tokens each generation request appends this step. Only populated under + # ragged verification; a uniform batch leaves it unused and the + # compressor keeps deriving the count from its next_n template bound. + self.gen_new_tokens_per_seq_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences,), + dtype=torch.int, + cache_name="gen_new_tokens_per_seq_cuda", + capture_graph=capture_graph, + ) + # Set per step by prepare(); None means "uniform, use next_n". + self.gen_new_tokens_per_seq: Optional[torch.Tensor] = None + # new_comp_kv_lens_cuda is the number of new compressed tokens for the requests self.new_comp_kv_lens_cuda = { compress_ratio: self.get_empty( @@ -673,7 +686,13 @@ def _init_cache_buffer_data_pointers(self): if is_compress_layer(extend_compress_ratios[layer_idx]) } - def prepare(self): + def _prepare_impl(self): + # Overrides DSAtrtllmAttentionMetadata._prepare_impl, NOT prepare(): + # the base prepare() wraps this whole body in the write-after-read + # guard for the pinned staging rewritten below (cached_token_lens_cpu, + # cu_seq_lens, and everything the called + # helpers stage). Overriding prepare() directly would silently opt this + # class out of that guard. assert self.kv_cache_manager is not None assert self.request_ids is not None @@ -729,10 +748,13 @@ def prepare(self): # --- Per-ratio metadata --- # 1) CPU-side: compute scalar metadata (num_total_compressed_tokens, etc.) # 2) CUDA-side: fill *_cuda buffers via prepare_compressed_kv_metadata() - num_gen_tokens_per_seq = ( - num_gen_tokens // self.num_generations if self.num_generations > 0 else 0 - ) - self.num_gen_tokens_per_seq = num_gen_tokens_per_seq + # Under ragged verification requests append different numbers of tokens, + # so the integer division below is meaningless. The compressor kernel + # still needs a compile-time loop bound, and its loops are already + # guarded, so the GLOBAL max window is the scalar to hand it (the + # device-window prologue may re-rank past this step's batch max) -- + # the exact per-request counts ride along in gen_new_tokens_per_seq. + num_gen_tokens_per_seq = self._sync_gen_tokens_per_seq(num_gen_tokens) num_contexts = self.num_contexts num_generations = self.num_generations kv_lens_slice = kv_lens[:num_requests] @@ -787,6 +809,46 @@ def prepare(self): self._compress_ratios_sorted, ) + def _sync_gen_tokens_per_seq(self, num_gen_tokens: int) -> int: + """Set ``num_gen_tokens_per_seq`` / ``gen_new_tokens_per_seq`` together. + + Uniform batches divide the flat generation-token count by the request + count. Ragged batches have no such quotient, so the scalar becomes the + GLOBAL max window -- ``NEXT_N`` is a template loop bound whose + iterations are already guarded, so any upper bound is correct, and it + must be the global one because the device-window prologue may re-rank + rows up to the config max AFTER this staging runs (this step's batch + max would truncate Phase 1/3 for such rows). The true per-request + counts travel beside it on the device. + + The vector is a D2D copy from the seq-lens staging, NOT an H2D from a + pinned host list: this helper is captured inside ``on_update_kv_lens``, + and a captured H2D node would re-load the host's shape split S from + the pinned buffer on every replay -- silently clobbering the device + windows w and making the compressor place its persistent paged + KV/score state at positions shifted by (w - S). ``_seq_lens_cuda`` + holds the same S at staging time and is rewritten to w by + ``apply_device_ragged_layout`` before replay, so sourcing from it is + exact for both host and device windows. + + Single writer on purpose: the scalar and the vector describe the same + thing, and every place that recomputed one without the other was a way + for them to disagree. + """ + if self.is_ragged_verify and self.num_generations > 0: + self.num_gen_tokens_per_seq = 1 + self.max_draft_tokens + n = self.num_generations + nc = self.num_contexts + gen_new_tokens = self.gen_new_tokens_per_seq_cuda[:n] + gen_new_tokens.copy_(self._seq_lens_cuda[nc:nc + n]) + self.gen_new_tokens_per_seq = gen_new_tokens + else: + self.num_gen_tokens_per_seq = ( + num_gen_tokens // self.num_generations if self.num_generations > 0 else 0 + ) + self.gen_new_tokens_per_seq = None + return self.num_gen_tokens_per_seq + def prepare_compressed_kv_metadata( self, kv_lens: torch.Tensor, @@ -862,9 +924,12 @@ def on_update_kv_lens(self): cached_tokens = kv_lens - seq_lens num_gen_tokens = num_tokens - self.num_ctx_tokens - self.num_gen_tokens_per_seq = ( - num_gen_tokens // self.num_generations if self.num_generations > 0 else 0 - ) + # Must go through the same helper as prepare(): a plain integer division + # here would silently undo the ragged batch maximum, and the compressor + # reads this value as its next_n loop bound. A bound below a request's + # real window truncates Phase 1 and the tail KV/score state is never + # written -- no error, just corrupted compression windows from then on. + self._sync_gen_tokens_per_seq(num_gen_tokens) # Reuse prepare()'s host-computed ctx sizes unless the extend_ctx path # (num_chunked_ctx_requests > 0) may have mutated ctx-row kv_lens on diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index a57fdadf0776..c5540472285e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -582,6 +582,29 @@ class IndexerPrefillChunkMetadata: k_token_end: int # K token end index in batch +@dataclass(frozen=True) +class TokenMajorGenView: + """One attention row per generation query token. + + Handed only to the MLA rope and sparse-MLA generation ops, which derive + their sequence count from ``host_context_lengths.size(0)``. Everything else + keeps seeing the batch-major metadata, where ``num_seqs`` still means + "number of requests". + """ + num_rows: int + sequence_length: torch.Tensor + host_past_key_value_lengths: torch.Tensor + host_context_lengths: torch.Tensor + prompt_lens_cuda: torch.Tensor + host_request_types: torch.Tensor + kv_cache_block_offsets: torch.Tensor + #: Static row ceiling, NOT this step's row count. The attention op caches + #: by (beam_width, max_num_requests, attention_window_size) and reserves its + #: semaphore array from it, so a per-step value would thrash that cache and + #: re-reserve on every shape change. + max_num_rows: int + + @dataclass(init=False) class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): """Attention metadata for DSA (Dense Sparse Attention) with indexer state.""" @@ -634,6 +657,94 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): # Number of compressed KV tokens for context requests num_ctx_kv_tokens: int = 0 gen_indexer_kv_lens_cuda_runtime: Optional[torch.Tensor] = None + # Ragged verification: query tokens per generation request this step, + # host-side and in batch order. None (the default) means the uniform + # `1 + max_draft_tokens` the strided expansions assume. Host-side because + # reading repeats off a device tensor would sync. + ragged_verify_lens: Optional[List[int]] = None + # Device-window mode (TLLM_DSPARK_DEVICE_WINDOWS): on ragged steps the + # host list above is only a SHAPE split summing to the bucket; the true + # per-request windows are chosen on device after prepare() and written + # through apply_device_ragged_layout. Host reads of window values must + # then be treated as bounds (see _prepare_impl). + device_windows_mode: bool = False + # Query tokens per generation request on a *uniform* batch this step + # (`1 + runtime_draft_len`). Zero means "not published"; expansions then + # fall back to the static `1 + max_draft_tokens`. + # + # Deliberately NOT `max_draft_tokens`: that ceiling sizes the expanded + # buffers and must never change under CUDA-graph capture, while the + # *stride* must follow the runtime tier -- `Indexer.forward` slices the + # same buffers by the true token count, so a ceiling stride makes every + # request past the first read a neighbour's rows. + runtime_tokens_per_gen_step: int = 0 + + @property + def is_ragged_verify(self) -> bool: + """Whether generation requests have per-request query lengths.""" + return self.ragged_verify_lens is not None + + @property + def gen_token_stride(self) -> int: + """Query tokens contributed by one generation request, uniform batches. + + The runtime tier when the engine published one, else the static + ceiling. Capacity questions must keep using `1 + max_draft_tokens`; + only the stride follows the tier. + """ + stride = int(self.runtime_tokens_per_gen_step) + if stride <= 0: + return 1 + self.max_draft_tokens + return stride + + def gen_token_repeat_list(self) -> List[int]: + """Per-generation-request query-token counts, on the host. + + Uniform batches get `gen_token_stride` per request; ragged batches get + the scheduler's windows. Host-side on purpose: the total feeds buffer + slicing, and reducing it off a device tensor would sync every step. + """ + if self.ragged_verify_lens is None: + return [self.gen_token_stride] * self.num_generations + return list(self.ragged_verify_lens) + + def gen_token_repeats(self, device=None) -> torch.Tensor: + """:meth:`gen_token_repeat_list` as a repeat vector for the given device.""" + return torch.tensor(self.gen_token_repeat_list(), + dtype=torch.long, + device=device) + + def expand_per_gen_token(self, + values: torch.Tensor, + dim: int = 0) -> Tuple[torch.Tensor, int]: + """Repeat `values` once per query token, returning `(expanded, num_tokens)`. + + Uniform batches take the scalar `repeat_interleave` (no repeats tensor, + no copy). Ragged batches must pass `output_size` (a tensor-repeats + interleave otherwise reads the cumsum back to the host) and stage the + repeats through pinned memory so the H2D copy is non-blocking. + """ + repeats = self.gen_token_repeat_list() + num_tokens = sum(repeats) + if self.ragged_verify_lens is None: + stride = self.gen_token_stride + return values.repeat_interleave(stride, dim=dim), num_tokens + # Device values take repeats from the PERSISTENT staged buffer, never + # a per-call pinned temporary: that call site runs inside the captured + # graph (via on_update_kv_lens), where a temporary would bake its host + # pointer into the captured H2D copy. The buffer is staged by + # prepare_for_spec_decode and overwritten in place by the + # device-window prologue. Host-tensor calls (prepare-time expansions + # of host kv lens / block tables) keep a plain host repeats tensor. + if values.device == self.gen_token_repeats_cuda.device: + repeats_dev = self.gen_token_repeats_cuda[:self.num_generations] + else: + repeats_dev = torch.tensor(repeats, + dtype=torch.long, + device=values.device) + return (values.repeat_interleave(repeats_dev, + dim=dim, + output_size=num_tokens), num_tokens) def __init__(self, *args, **kwargs): """Initialize DSA metadata with SM count and indexer chunk size.""" @@ -709,6 +820,24 @@ def __post_init__(self): self.create_buffers_for_indexer(capture_graph=capture_graph) def prepare(self): + # WAR guard over EVERYTHING prepare() stages through pinned memory: + # under the overlap scheduler last step's non_blocking H2D copies can + # still be queued when this rewrite starts, so wait on the event + # recorded after the previous prepare's copies before touching any. + if not torch.cuda.is_current_stream_capturing(): + evt = getattr(self, "_prepare_stage_evt", None) + if evt is None: + self._prepare_stage_evt = torch.cuda.Event() + else: + evt.synchronize() + try: + self._prepare_impl() + finally: + evt = getattr(self, "_prepare_stage_evt", None) + if evt is not None and not torch.cuda.is_current_stream_capturing(): + evt.record() + + def _prepare_impl(self): super().prepare() self._invalidate_pool_view_cache() # Cross-layer indexer sharing is per-step state; clear it so a "shared" @@ -732,6 +861,20 @@ def prepare(self): else: kv_lens = cached_token_lens + self.seq_lens_kv + if self.device_windows_mode and self.is_ragged_verify: + # Device-window steps: the host's per-request windows are only a + # SHAPE split; the true windows are chosen on device after + # prepare(). Every host-side read of kv_lens below is a bound or + # is rebuilt in-graph from kv_lens_cuda (on_update_kv_lens), so + # inflate the gen entries to the full drafted block -- a bound no + # device-chosen window can exceed. kv_lens_cuda keeps the exact + # shape-split values (staged by super().prepare() above); the + # prologue's per-request delta corrects it to the true windows. + nc, ns = self.num_contexts, self.num_seqs + kv_lens = kv_lens.clone() + kv_lens[nc:ns] += (1 + self.max_draft_tokens) - self.seq_lens_kv[nc:ns] + self.kv_lens[nc:ns] = kv_lens[nc:ns] + # For mla_rope_append_paged_kv_assign_q self.prepare_for_mla_rope_append(cached_token_lens, kv_lens) @@ -828,6 +971,10 @@ def on_update_kv_lens(self): # Per-step state for cross-layer indexer sharing; clear at the step # boundary so a "shared" layer never reuses a stale top-k. self.shared_topk_indices = None + # kv_lens_cuda may have just been corrected; the ragged per-row extent + # is derived from it and has to follow. + self.refresh_ragged_row_kv_lens() + self.refresh_token_major_gen_rows() if self.kv_cache_manager is not None and self.num_tokens > 0: seq_lens = self.seq_lens_cuda[:self.num_seqs] @@ -894,13 +1041,11 @@ def on_update_kv_lens(self): self.scheduler_metadata_buffer_full_next_n.copy_( scheduler_metadata_buffer_full_next_n, non_blocking=True) if self.use_expanded_buffers_for_mtp: - num_draft_tokens = 1 + self.max_draft_tokens - num_tokens = self.num_generations * num_draft_tokens - kv_lens_expanded = torch.stack([gen_indexer_kv_lens] * - num_draft_tokens, - dim=0) - self.kv_lens_expanded_cuda[:num_tokens] = \ - kv_lens_expanded.transpose(0, 1).contiguous().flatten() + # One row per query token; per-request repeats under ragged + # verification (see ``expand_per_gen_token``). + expanded, num_tokens = self.expand_per_gen_token( + gen_indexer_kv_lens) + self.kv_lens_expanded_cuda[:num_tokens] = expanded scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), _DG_SCHEDULE_BLOCK_KV, self.num_sms) @@ -938,7 +1083,10 @@ def _compute_kv_lens_row_reorder(self): # num_generations alone is only correct for next_n == 2; for next_n == 1 # it engages inside the measured regression band, and for next_n == 4 it # misses the win region between 2*num_sms and 4*num_sms rows. - next_n = 1 + self.max_draft_tokens + # Runtime stride, not the static ceiling: this is a row-count threshold, + # and at a shorter verify tier the static value overstates the rows by + # the tier ratio, engaging GVR outside its measured win region. + next_n = self.gen_token_stride if (self.enable_heuristic_topk and self.use_cute_dsl_topk and self.num_generations * next_n >= 2 * self.num_sms): gen_kv_lens = self.kv_lens_cuda[self.num_contexts:self.num_seqs] @@ -1106,6 +1254,23 @@ def create_buffers_for_indexer(self, capture_graph=False): ) self.host_req_idx_per_token = torch.empty_like( self.req_idx_per_token, device='cpu', pin_memory=prefer_pinned()) + # Per-gen-request query-token counts as a PERSISTENT device tensor. + # The in-graph ragged expansion (``on_update_kv_lens`` -> + # ``expand_per_gen_token``) must read repeats from a stable address: + # a per-call pinned temporary would bake the temporary's host pointer + # into the captured H2D. Staged in prepare_for_spec_decode; the + # device-window prologue overwrites it in place. + self.gen_token_repeats_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences, ), + cache_name="gen_token_repeats_cuda", + dtype=torch.int64, + capture_graph=capture_graph, + ) + self.host_gen_token_repeats = torch.empty_like( + self.gen_token_repeats_cuda, + device='cpu', + pin_memory=prefer_pinned()) # Block table for topk_indices conversion (shared for context and generation) self.block_table = self.get_empty( self.cuda_graph_buffers, @@ -1200,7 +1365,7 @@ def create_buffers_for_indexer(self, capture_graph=False): # path does not, so skip the allocation when use_cute_dsl_topk. if not self.use_cute_dsl_topk: max_gen_tokens = self.max_num_sequences * ( - 1 + self.max_draft_tokens) + 1 + self._draft_sizing_cap) self.heuristic_scratch_values = self.get_empty( self.cuda_graph_buffers, (max_gen_tokens, self.num_sparse_topk), @@ -1229,6 +1394,23 @@ def create_buffers_for_indexer(self, capture_graph=False): # Create expanded buffers for MTP support self.create_expanded_buffers(capture_graph=capture_graph) + @property + def _draft_sizing_cap(self) -> int: + """Static draft-token ceiling for buffer SIZING (not per-step logic). + + Every expanded buffer must be shaped from one process-wide ceiling, + never the per-step ``max_draft_tokens``: the pooled backing memory is + shared between the base metadata and every CUDA-graph clone, and two + instances viewing the same bytes with different row counts corrupt + each other. ``update_spec_dec_param`` records the config-wide + ``max_total_draft_tokens`` here; before the first call the mutable + value is all there is. + """ + cap = getattr(self, "_draft_alloc_cap", None) + if cap is None: + return self.max_draft_tokens + return max(cap, self.max_draft_tokens) + def _create_kv_lens_2d_buffer(self, capture_graph=False): """Pre-allocated buffer for the DeepGEMM 2D context_lens API. @@ -1238,7 +1420,7 @@ def _create_kv_lens_2d_buffer(self, capture_graph=False): """ self.kv_lens_cuda_2d = self.get_empty( self.cuda_graph_buffers, - (self.max_num_sequences, 1 + self.max_draft_tokens), + (self.max_num_sequences, 1 + self._draft_sizing_cap), cache_name="kv_lens_cuda_2d", dtype=torch.int32, capture_graph=capture_graph, @@ -1249,7 +1431,7 @@ def create_expanded_buffers(self, capture_graph=False): """Create expanded KV-length and block-table buffers for speculative decoding.""" self.kv_lens_expanded_cuda = self.get_empty( self.cuda_graph_buffers, - (self.max_num_sequences * (1 + self.max_draft_tokens), ), + (self.max_num_sequences * (1 + self._draft_sizing_cap), ), cache_name="kv_lens_expanded_cuda", dtype=torch.int32, capture_graph=capture_graph, @@ -1259,10 +1441,159 @@ def create_expanded_buffers(self, capture_graph=False): device='cpu', pin_memory=prefer_pinned(), ) + # Per-row causal KV extent for the decode top-k, in *uncompressed* token + # space (the kernel divides by compress_ratio itself). Only populated on + # ragged batches: the uniform kernel reconstructs the same quantity from + # next_n, and passing nullptr keeps that path bit-identical. + # + # Distinct from kv_lens_expanded_cuda, which carries indexer (possibly + # compressed) lengths for the paged-MQA-logits call. The top-k kernel is + # handed kv_lens_cuda_runtime instead, so mixing the two would apply the + # compression divisor twice. + self.row_kv_lens_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences * (1 + self._draft_sizing_cap), ), + cache_name="row_kv_lens_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.row_kv_lens_host = torch.zeros_like( + self.row_kv_lens_cuda, + device='cpu', + pin_memory=prefer_pinned(), + ) + # The per-row offset `-verify_len + position + 1`, kept separately so + # the extent can be rebuilt from a corrected kv_lens without going back + # to the host. It depends only on the verify windows, which do not move + # after prepare(); the kv lengths do. + self.row_kv_correction_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences * (1 + self._draft_sizing_cap), ), + cache_name="row_kv_correction_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.row_kv_correction_host = torch.zeros_like( + self.row_kv_correction_cuda, + device='cpu', + pin_memory=prefer_pinned(), + ) + # Which request each query row belongs to. Persistent because the + # refresh below runs inside the CUDA-graph capture body: building this + # per call would record a memcpy from a Python temporary that is freed + # before the first replay. + self.row_req_idx_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences * (1 + self._draft_sizing_cap), ), + cache_name="row_req_idx_cuda", + dtype=torch.long, + capture_graph=capture_graph, + ) + self.row_req_idx_host = torch.zeros_like( + self.row_req_idx_cuda, + device='cpu', + pin_memory=prefer_pinned(), + ) + # Rows covered by the current step's ragged layout; 0 on uniform steps. + self._ragged_num_rows = 0 + + # ------------------------------------------------------------------ + # Token-major presentation of the generation half, for the MLA rope + # and sparse-MLA generation ops. Both derive their sequence count from + # `host_context_lengths.size(0)` and require the query tokens to + # divide evenly across it; a ragged batch does not, so hand them one + # row per generation token (seq_len == 1). + # + # These are a PARALLEL set of views, never the metadata's own fields: + # `num_seqs` is read as a request count throughout and must not be + # redefined. Rows are prefix-preserving -- context rows first, + # unchanged, then one row per generation token -- keeping + # `num_contexts`, `num_ctx_tokens` and the ops' `seq_offset` correct. + # + # Capacity is the static ceiling and must never be recomputed from a + # runtime count: these live in the process-wide buffer pool keyed by + # `cache_name`, so reallocation would move addresses baked into + # captured graphs. `2 + max_draft_tokens` per sequence covers one + # context row plus the widest generation window. + attn_rows_cap = self.max_num_sequences * (2 + self._draft_sizing_cap) + self._attn_num_rows = 0 + self.attn_row_kv_lens_cuda = self.get_empty( + self.cuda_graph_buffers, + (attn_rows_cap, ), + cache_name="attn_row_kv_lens_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + # The per-row offset from its request's kv_len, kept separately so the + # extent can be rebuilt on device from a corrected kv_lens without going + # back to the host -- `_preprocess_inputs` rewrites kv_lens_cuda inside + # the capture body under the overlap scheduler. + self.attn_row_kv_correction_cuda = self.get_empty( + self.cuda_graph_buffers, + (attn_rows_cap, ), + cache_name="attn_row_kv_correction_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.attn_row_req_idx_cuda = self.get_empty( + self.cuda_graph_buffers, + (attn_rows_cap, ), + cache_name="attn_row_req_idx_cuda", + dtype=torch.long, + capture_graph=capture_graph, + ) + self.attn_row_kv_lens_host = torch.zeros( + (attn_rows_cap, ), dtype=torch.int32, + pin_memory=prefer_pinned()) + self.attn_row_kv_correction_host = torch.zeros( + (attn_rows_cap, ), dtype=torch.int32, + pin_memory=prefer_pinned()) + self.attn_row_req_idx_host = torch.zeros( + (attn_rows_cap, ), dtype=torch.long, pin_memory=prefer_pinned()) + # host_context_lengths: the shape oracle both ops key off. The op + # reads request_types over its whole num_seqs -- the ROW count under + # this presentation -- so the array must be row-sized. Generation is + # 1, context 0; prefill the whole capacity as generation so a step + # only stamps its context prefix. + self.attn_row_request_types_host = torch.ones( + (attn_rows_cap, ), dtype=torch.int32, + pin_memory=prefer_pinned()) + # Device twin of attn_row_prompt_lens_cpu: the forward path asserts + # every runtime view shares batch_size with kv_lens_cuda_runtime. + self.attn_row_prompt_lens_cuda = self.get_empty( + self.cuda_graph_buffers, + (attn_rows_cap, ), + cache_name="attn_row_prompt_lens_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.attn_row_prompt_lens_cpu = torch.zeros( + (attn_rows_cap, ), dtype=torch.int32, + pin_memory=prefer_pinned()) + # The main KV block table, one row per query token. Distinct from + # `block_table_expanded` below, which carries *indexer* pool indices at + # the indexer's own compression. Same [num_pools, num_seqs, 2, + # max_blocks_per_seq] layout the attention op reads, expanded on the + # sequence axis only: it substitutes for `kv_cache_block_offsets` + # wholesale, so rank and axis order must match. + num_attention_op_pools = getattr(self.kv_cache_manager, + "num_attention_op_pools", + self.kv_cache_manager.num_pools) + self.attn_row_block_offsets = self.get_empty( + self.cuda_graph_buffers, + [ + num_attention_op_pools, attn_rows_cap, 2, + self.kv_cache_manager.max_blocks_per_seq + ], + cache_name="attn_row_block_offsets", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.block_table_expanded = self.get_empty( self.cuda_graph_buffers, [ - self.max_num_sequences * (1 + self.max_draft_tokens), + self.max_num_sequences * (1 + self._draft_sizing_cap), self.kv_cache_manager.max_blocks_per_seq ], cache_name="block_table_expanded", @@ -1309,18 +1640,29 @@ def update_spec_dec_param( spec_tree_manager, num_contexts=num_contexts) self.max_draft_tokens = max_draft_len + # The config-wide ceiling, identical every call; sizing must use this + # so base and every clone share one buffer layout (see + # _draft_sizing_cap). + self._draft_alloc_cap = max( + int(max_total_draft_tokens), + getattr(self, "_draft_alloc_cap", 0) or 0) capture_graph = self.is_cuda_graph - if self.kv_lens_cuda_2d.shape[1] != 1 + self.max_draft_tokens: + # GROW-ONLY ratchet, never an equality check: captured graphs bake the + # buffer addresses in, so re-creating buffers when max_draft_len moves + # per step corrupts every captured graph. Consumers slice + # [:rows]/[:width]. + prev_cap = getattr(self, "_expanded_alloc_draft_cap", -1) + ratchet = max(self._draft_sizing_cap, prev_cap) + if ratchet > prev_cap: + self._expanded_alloc_draft_cap = ratchet self._create_kv_lens_2d_buffer(capture_graph=capture_graph) - init_shape = self.kv_lens_expanded_host.shape[0] - if self.max_num_sequences * (1 + self.max_draft_tokens) != init_shape: self.create_expanded_buffers(capture_graph=capture_graph) # Resize heuristic scratch buffer for new max_draft_tokens. # Skip when use_cute_dsl_topk (GVR path never consumes it), matching # the allocation guard in create_buffers_for_indexer. if self.enable_heuristic_topk and not self.use_cute_dsl_topk: max_gen_tokens = self.max_num_sequences * ( - 1 + self.max_draft_tokens) + 1 + self._draft_sizing_cap) self.heuristic_scratch_values = self.get_empty( self.cuda_graph_buffers, (max_gen_tokens, self.num_sparse_topk), @@ -1328,12 +1670,10 @@ def update_spec_dec_param( dtype=torch.float32, capture_graph=capture_graph, ) - # The Radix-split-work scratch (radix_aux_*) is sized the same way - # (num_seqs * (1 + max_draft_tokens) rows) and is allocated - # unconditionally, so it must be resized here too -- otherwise the - # cpp indexer_topk_decode op aborts once MTP raises max_draft_tokens - # ("radix_aux_* must hold at least num_rows*blocks_per_row*index_topk - # elements"). + # radix_aux_* is sized by the same generation-row count and + # allocated unconditionally, so it must be resized here too or the + # cpp indexer_topk_decode op aborts once MTP raises + # max_draft_tokens. self._create_radix_aux_buffers(capture_graph=capture_graph) def _get_pool_block_indices(self) -> torch.Tensor: @@ -1475,20 +1815,52 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): # - No distinction between sm90 and sm100 is needed once MTP3 is supported on sm90. # - Remove this once fp8_paged_mqa_logits supports an arbitrary number of MTP draft tokens. use_dsl = self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits - self.use_expanded_buffers_for_mtp = (not use_dsl and ( - (self.max_draft_tokens > 1 and get_sm_version() == 90) or - ((self.max_draft_tokens == 2 or self.max_draft_tokens > 3) - and get_sm_version() >= 100))) + if not self.is_ragged_verify: + # A uniform step must not inherit a previous ragged step's row + # count: refresh_ragged_row_kv_lens and ragged_row_kv_lens key + # off it. + self._ragged_num_rows = 0 + if self.is_ragged_verify: + # The DSL paged-MQA path is structurally uniform (one broadcast + # kv_len per request, factor * atom == next_n); route ragged + # batches to the expanded one-row-per-token path instead. + use_dsl = False + # The dense XQA spec-dec metadata (spec_decoding_generation_lengths, + # packed mask) is filled with a single scalar window and cannot + # describe per-request verify lengths; assert it stays off so a + # future change fails here rather than silently verifying every + # request at the batch-wide window. + assert not self.is_spec_decoding_enabled, ( + "ragged verification requires the dense spec-decoding metadata " + "to be off: spec_decoding_generation_lengths and the packed " + "mask are filled with a single scalar window, which cannot " + "describe per-request verify lengths") + # Keyed on the static ceiling, not this step's tier, so the layout + # regime never changes per step; the expanded path is correct at every + # tier and required for ragged batches anyway. + self.use_expanded_buffers_for_mtp = ( + not use_dsl + and (self.is_ragged_verify or + (self.max_draft_tokens > 1 and get_sm_version() == 90) or + ((self.max_draft_tokens == 2 or self.max_draft_tokens > 3) + and get_sm_version() >= 100))) if self.use_expanded_buffers_for_mtp: - # Expand kv_lens_cuda (only generation) - num_tokens = self.num_generations * (1 + self.max_draft_tokens) + if self.is_ragged_verify: + # Stage the persistent device repeats the in-graph expansion + # reads (see expand_per_gen_token); must land before the + # expansions below and before any capture of this step. + n_gen = self.num_generations + self.host_gen_token_repeats[:n_gen].copy_( + torch.tensor(self.gen_token_repeat_list(), + dtype=torch.int64)) + self.gen_token_repeats_cuda[:n_gen].copy_( + self.host_gen_token_repeats[:n_gen], non_blocking=True) + # Expand kv_lens_cuda (only generation): one row per query token, + # repeating each request's kv_len as many times as it has tokens. gen_kv_lens = self.get_indexer_kv_lens( kv_lens[self.num_contexts:self.num_seqs]) - gen_kv_lens_expanded = torch.stack([gen_kv_lens] * - (1 + self.max_draft_tokens), - dim=0) - gen_kv_lens_expanded = gen_kv_lens_expanded.transpose( - 0, 1).contiguous().flatten() + gen_kv_lens_expanded, num_tokens = self.expand_per_gen_token( + gen_kv_lens) self.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded) self.kv_lens_expanded_cuda[:num_tokens].copy_( self.kv_lens_expanded_host[:num_tokens], non_blocking=True) @@ -1500,8 +1872,8 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): max_len = self.host_indexer_k_cache_block_offsets.shape[1] gen_block_tensor = self.host_indexer_k_cache_block_offsets[ self.num_contexts:self.num_seqs, :max_len] - expanded_blocks = gen_block_tensor.repeat_interleave( - 1 + self.max_draft_tokens, dim=0) + expanded_blocks, _ = self.expand_per_gen_token(gen_block_tensor, + dim=0) self.host_block_table_expanded[:num_tokens, :max_len].copy_( expanded_blocks, non_blocking=True) self.block_table_expanded[:num_tokens].copy_( @@ -1509,10 +1881,17 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): non_blocking=True) self.block_table_expanded.clamp_(min=0) + if self.is_ragged_verify: + self._prepare_ragged_row_kv_lens(kv_lens) + self._prepare_token_major_gen_rows(kv_lens) + self.expand_for_dsl = (use_dsl and self.kv_cache_manager is not None and self.max_draft_tokens >= 1) if self.expand_for_dsl and self.num_generations > 0: - next_n = 1 + self.max_draft_tokens + # `Indexer.forward` validates the cached split against the runtime + # next_n (`factor * atom == next_n`), so the picker has to be fed the + # same value or the split is silently declined on every trimmed step. + next_n = self.gen_token_stride kernel_atoms = (1, 2, 3) if self.kv_cache_manager.use_fp4 else (1, 2, 3, 4) @@ -1550,7 +1929,262 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): self.block_table_expanded.clamp_(min=0) else: self.dsl_expand_factor = 1 - self.dsl_atom = 1 + self.max_draft_tokens + self.dsl_atom = self.gen_token_stride + + def _prepare_ragged_row_kv_lens(self, kv_lens: torch.Tensor) -> None: + """Populate ``row_kv_lens_cuda``: how far back each query row may attend. + + The uniform kernel reconstructs this from ``next_n``, which is + meaningless on a ragged batch, so materialize it per row: row ``o`` of + a request verifying ``v`` positions attends to ``kv_len - v + o + 1`` + tokens (``v = next_n`` recovers the uniform formula, so that path + keeps passing ``nullptr``). Assembled on the host -- ``kv_lens`` is + already a CPU tensor -- and copied once through pinned memory; + building it on device would sync. + """ + verify_lens = self.ragged_verify_lens + if not verify_lens: + return + gen_kv_lens = kv_lens[self.num_contexts:self.num_seqs].tolist() + assert len(gen_kv_lens) == len(verify_lens), ( + f"ragged verify lengths {len(verify_lens)} != generation requests " + f"{len(gen_kv_lens)}") + # `ragged_verify_lens` is published separately from this step's token + # layout; check they describe the same step. + expected_gen_tokens = self.num_tokens - self.num_ctx_tokens + assert sum(verify_lens) == expected_gen_tokens, ( + f"ragged verify windows sum to {sum(verify_lens)} but this step's " + f"token layout has {expected_gen_tokens} generation tokens; the " + f"windows belong to a different step than seq_lens does") + + rows: List[int] = [] + for kv_len, verify_len in zip(gen_kv_lens, verify_lens): + base = int(kv_len) - int(verify_len) + 1 + rows.extend(range(base, base + int(verify_len))) + + corrections: List[int] = [] + req_idx: List[int] = [] + for request_pos, verify_len in enumerate(verify_lens): + corrections.extend(range(1 - int(verify_len), 1)) + req_idx.extend([request_pos] * int(verify_len)) + + num_tokens = len(rows) + self._ragged_num_rows = num_tokens + # The write-after-read hazard on these pinned buffers (rewritten each + # ragged step, copied non_blocking under the overlap scheduler) is + # fenced by the event guard in prepare(). + self.row_kv_lens_host[:num_tokens].copy_( + torch.tensor(rows, dtype=torch.int32)) + self.row_kv_lens_cuda[:num_tokens].copy_( + self.row_kv_lens_host[:num_tokens], non_blocking=True) + self.row_kv_correction_host[:num_tokens].copy_( + torch.tensor(corrections, dtype=torch.int32)) + self.row_req_idx_host[:num_tokens].copy_( + torch.tensor(req_idx, dtype=torch.long)) + self.row_kv_correction_cuda[:num_tokens].copy_( + self.row_kv_correction_host[:num_tokens], non_blocking=True) + self.row_req_idx_cuda[:num_tokens].copy_( + self.row_req_idx_host[:num_tokens], non_blocking=True) + + def refresh_ragged_row_kv_lens(self) -> None: + """Rebuild the per-row extent from the corrected ``kv_lens_cuda``. + + Under the overlap scheduler ``_preprocess_inputs`` corrects + ``kv_lens_cuda`` after prepare(); the extent must follow or the top-k + kernel's arguments disagree. Device-only: one gather plus one add, no + sync. + """ + num_tokens = self._ragged_num_rows + if not self.is_ragged_verify or num_tokens <= 0: + return + gen_kv_lens = self.kv_lens_cuda[self.num_contexts:self.num_seqs] + # Gather, not repeat_interleave: this runs inside the capture body, + # where a tensor-repeats interleave would allocate and sync. The map + # and correction were staged at prepare() time. + row_kv_lens = self.row_kv_lens_cuda[:num_tokens] + torch.index_select(gen_kv_lens, 0, + self.row_req_idx_cuda[:num_tokens], + out=row_kv_lens) + row_kv_lens.add_(self.row_kv_correction_cuda[:num_tokens]) + + def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: + """Build the one-row-per-generation-token view for the MLA gen ops. + + See the buffer comments in :meth:`create_expanded_buffers`. Rows are + ``[contexts unchanged] ++ [one per generation token]``; row ``o`` of a + request verifying ``v`` positions attends to ``kv_len - v + o + 1`` + tokens. Assembled on the host and copied once through pinned memory; + building it on device would sync. + """ + verify_lens = self.ragged_verify_lens + ngt = self._ragged_num_rows + if not verify_lens or ngt <= 0: + self._attn_num_rows = 0 + return + nc = self.num_contexts + rows = nc + ngt + cap = self.attn_row_kv_lens_cuda.shape[0] + assert rows <= cap, ( + f"token-major generation needs {rows} rows but the buffers hold " + f"{cap}; capacity is the static ceiling and must not be resized " + f"after any CUDA graph was captured") + + # Context rows pass through untouched, so the ops' seq_offset still + # points at the first generation row. + if nc: + ctx_kv = kv_lens[:nc].to(torch.int32) + self.attn_row_kv_lens_host[:nc].copy_(ctx_kv) + self.attn_row_kv_correction_host[:nc].zero_() + self.attn_row_req_idx_host[:nc].copy_( + torch.arange(nc, dtype=torch.long)) + self.attn_row_prompt_lens_cpu[:nc].copy_( + self.prompt_lens_cpu[:nc].to(torch.int32)) + + # Reuse the per-row extents/corrections/request indices + # _prepare_ragged_row_kv_lens just built. The pinned-buffer WAR hazard + # is fenced by the event guard in prepare(). + self.attn_row_kv_lens_host[nc:rows].copy_( + self.row_kv_lens_host[:ngt]) + self.attn_row_kv_correction_host[nc:rows].copy_( + self.row_kv_correction_host[:ngt]) + # row_req_idx is batch-relative among generation requests; the attention + # rows are indexed over the whole batch, so shift past the contexts. + torch.add(self.row_req_idx_host[:ngt], nc, + out=self.attn_row_req_idx_host[nc:rows]) + # One query token per generation row; contexts keep their real q + # length. + self.attn_row_prompt_lens_cpu[nc:rows].fill_(1) + # Capacity is prefilled with kGENERATION, so only the context prefix + # needs stamping, and nc is small. + if nc: + self.attn_row_request_types_host[:nc].zero_() + self.attn_row_request_types_host[nc:rows].fill_(1) + + self.attn_row_kv_lens_cuda[:rows].copy_( + self.attn_row_kv_lens_host[:rows], non_blocking=True) + self.attn_row_kv_correction_cuda[:rows].copy_( + self.attn_row_kv_correction_host[:rows], non_blocking=True) + self.attn_row_req_idx_cuda[:rows].copy_( + self.attn_row_req_idx_host[:rows], non_blocking=True) + self.attn_row_prompt_lens_cuda[:rows].copy_( + self.attn_row_prompt_lens_cpu[:rows], non_blocking=True) + + self._attn_num_rows = rows + self.refresh_token_major_block_table() + + def refresh_token_major_gen_rows(self) -> None: + """Rebuild the per-row KV extent from the corrected ``kv_lens_cuda``. + + Mirrors :meth:`refresh_ragged_row_kv_lens`: the overlap scheduler + corrects ``kv_lens_cuda`` on device after prepare(), inside the + capture body. One gather plus one add; no allocation, no D2H. + """ + rows = self._attn_num_rows + if not self.is_ragged_verify or rows <= 0: + return + torch.index_select(self.kv_lens_cuda[:self.num_seqs], 0, + self.attn_row_req_idx_cuda[:rows], + out=self.attn_row_kv_lens_cuda[:rows]) + self.attn_row_kv_lens_cuda[:rows].add_( + self.attn_row_kv_correction_cuda[:rows]) + + def refresh_token_major_block_table(self) -> None: + """Expand the main KV block table to one row per query token, on device. + + A verify window never spans a request, so this is a gather by the + row->request map on the already-resident table. Not refreshed from + ``on_update_kv_lens``: block IDs do not move within a step. + """ + rows = self._attn_num_rows + if rows <= 0 or self.kv_cache_block_offsets is None: + return + src = self.kv_cache_block_offsets + # Gather along the sequence axis (1); the pool and K/V axes pass + # through unchanged. + pools = min(src.shape[0], self.attn_row_block_offsets.shape[0]) + width = min(src.shape[-1], self.attn_row_block_offsets.shape[-1]) + torch.index_select( + src[:pools, :, :, :width], 1, self.attn_row_req_idx_cuda[:rows], + out=self.attn_row_block_offsets[:pools, :rows, :, :width]) + + def token_major_gen_view(self) -> Optional["TokenMajorGenView"]: + """The token-major presentation for this step, or None. + + Must be None on every uniform batch: at ``s_q == 1`` trtllm-gen + selects different kernels, so presenting a uniform batch token-major + would not stay bit-identical. + """ + rows = self._attn_num_rows + if not self.is_ragged_verify or rows <= 0: + return None + return TokenMajorGenView( + num_rows=rows, + sequence_length=self.attn_row_kv_lens_cuda[:rows], + host_past_key_value_lengths=self.attn_row_kv_lens_host[:rows], + host_context_lengths=self.attn_row_prompt_lens_cpu[:rows], + prompt_lens_cuda=self.attn_row_prompt_lens_cuda[:rows], + host_request_types=self.attn_row_request_types_host[:rows], + kv_cache_block_offsets=self.attn_row_block_offsets[:, :rows], + max_num_rows=self.attn_row_kv_lens_cuda.shape[0], + ) + + def ragged_row_kv_lens(self, num_tokens: int) -> Optional[torch.Tensor]: + """The per-row extent slice for this step, or None on a uniform batch.""" + if not self.is_ragged_verify: + return None + assert num_tokens == self._ragged_num_rows, ( + f"ragged row count moved between prepare ({self._ragged_num_rows}) " + f"and forward ({num_tokens}); the per-row extents would be read " + f"past what was populated") + return self.row_kv_lens_cuda[:num_tokens] + + def apply_device_ragged_layout(self, verify_lens: torch.Tensor, + req_idx: torch.Tensor, + kv_correction: torch.Tensor) -> None: + """Rebind this step's ragged layout to device-chosen windows. + + Pre-replay device ops only -- no host sync, stream-ordered before the + graph replay so every captured reader sees the new values. Overwrites + the same persistent buffers prepare() staged from the host shape + split; the per-row extents and every expanded/indexer artifact then + rebuild in-graph from them (``on_update_kv_lens``). + + Preconditions (the engine's prologue guarantees them): gen-only batch + (``num_contexts == 0``), ``verify_lens`` sums to the staged bucket + (``_ragged_num_rows``), and ``kv_lens_cuda`` has already been + delta-corrected from the shape split to these windows. + + Args: + verify_lens: ``[num_generations]`` int32 token windows. + req_idx: ``[bucket]`` int64 owning-row per packed token. + kv_correction: ``[bucket]`` int32, same ``1 - v .. 0`` convention + as the host staging (extent = kv_len + correction). + """ + nc = self.num_contexts + rows = self._ragged_num_rows + n_gen = self.num_generations + # Request-major: the true query lengths. Everything on_update_kv_lens + # derives per token (slot mappings, indptrs, dense top-k) follows. + self.seq_lens_cuda[nc:self.num_seqs] = verify_lens.to( + self.seq_lens_cuda.dtype) + self.gen_token_repeats_cuda[:n_gen] = verify_lens.to(torch.int64) + # Token-major maps, both the indexer's rows and the attention view's. + self.row_req_idx_cuda[:rows] = req_idx + self.row_kv_correction_cuda[:rows] = kv_correction + self.attn_row_req_idx_cuda[:rows] = req_idx + nc + self.attn_row_kv_correction_cuda[:rows] = kv_correction + self.req_idx_per_token[nc:nc + rows] = (req_idx + nc).to( + self.req_idx_per_token.dtype) + # Extents (row_kv_lens/attn_row_kv_lens) are rebuilt in-graph by the + # refresh functions; the block-table expansions are gathered per row + # here because block IDs do not move within a step and their refresh + # is prepare-time only. + self.refresh_token_major_block_table() + width = min(self.indexer_k_cache_block_offsets.shape[-1], + self.block_table_expanded.shape[-1]) + self.block_table_expanded[:rows, :width] = ( + self.indexer_k_cache_block_offsets[nc:self.num_seqs, :width] + .index_select(0, req_idx).clamp_(min=0)) def prepare_for_indexer_k_cache(self): # Build indexer_k_cache_block_offsets using pool block indices derived @@ -2234,7 +2868,12 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): metadata.kv_lens_cuda_runtime[num_contexts:num_contexts + num_generations]) metadata.gen_indexer_kv_lens_cuda_runtime = gen_seq_lens - next_n_cap = metadata.kv_lens_cuda_2d.shape[1] + # The 2D DeepGEMM metadata API encodes next_n into the schedule; + # build it at the *runtime* next_n `Indexer.forward` launches + # with, or a shorter verify tier describes more query positions + # than the kernel is given. + next_n_cap = min(metadata.gen_token_stride, + metadata.kv_lens_cuda_2d.shape[1]) metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap].copy_( gen_seq_lens.unsqueeze(-1).expand(-1, next_n_cap)) scheduler_metadata_buffer = get_paged_mqa_logits_metadata( @@ -2250,9 +2889,11 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): scheduler_metadata_buffer_full_next_n, non_blocking=True) else: # Expand schedule metadata buffer (only generation). The DeepGEMM - # API requires 2D; each expanded token becomes a (1,) row. - num_tokens = metadata.num_generations * (1 + - metadata.max_draft_tokens) + # API requires 2D; each expanded token becomes a (1,) row. The row + # count must be what the expansions actually populated + # (`gen_token_repeat_list`); a mismatch makes the kernel wait on + # tiles nobody produces (hang, not an error). + num_tokens = sum(metadata.gen_token_repeat_list()) scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( metadata.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), _DG_SCHEDULE_BLOCK_KV, metadata.num_sms) @@ -2709,9 +3350,12 @@ def sparse_attn_indexer( # Get decode lengths per request (from seq_lens) for validation gen_seq_lens = metadata.seq_lens[num_contexts:num_contexts + num_generations] - max_decode_len = gen_seq_lens.max().item() - min_decode_len = gen_seq_lens.min().item() - assert max_decode_len == min_decode_len, "max_decode_len != min_decode_len, we need padding" + if not metadata.is_ragged_verify: + # Uniformity only holds on the strided paths; ragged batches + # take the expanded (one row per token) path. + max_decode_len = gen_seq_lens.max().item() + min_decode_len = gen_seq_lens.min().item() + assert max_decode_len == min_decode_len, "max_decode_len != min_decode_len, we need padding" # Reshape q for decode phase: [num_gen_tokens, ...] -> [batch_size, next_n, ...] q_decode = q_fp8[num_ctx_tokens:num_ctx_tokens + num_gen_tokens, @@ -2721,7 +3365,14 @@ def sparse_attn_indexer( # Because fp8_paged_mqa_logits can only support next_n == 1/2/4 on sm100, and # next_n == 1/2 on sm90, for other next_n, we need to flatten the q_decode tensor # and expand the corresponding metadata. - if not metadata.use_expanded_buffers_for_mtp or next_n == 1: + # Ragged batches have no single next_n, so they always take the + # expanded (one row per token) branch. + if metadata.is_ragged_verify: + take_strided_path = False + else: + take_strided_path = (not metadata.use_expanded_buffers_for_mtp + or next_n == 1) + if take_strided_path: q_decode = q_decode.view(num_generations, -1, *q_fp8.shape[1:]) # 2D context_lens slice from the pre-allocated buffer; matches # q_decode's (batch, next_n) layout required by the new @@ -2750,7 +3401,14 @@ def sparse_attn_indexer( block_table = metadata.block_table_expanded[:num_tokens] scheduler_metadata_buffer = metadata.scheduler_metadata_buffer_expanded - assert num_gen_tokens == batch_size * next_n + # The expanded buffers must cover exactly the tokens handed to the + # kernel. + if metadata.is_ragged_verify: + assert num_gen_tokens == sum(metadata.ragged_verify_lens), ( + f"ragged gen tokens {num_gen_tokens} != " + f"sum(verify_lens) {sum(metadata.ragged_verify_lens)}") + else: + assert num_gen_tokens == batch_size * next_n weights_decode = weights[num_ctx_tokens:num_ctx_tokens + num_gen_tokens, ...] @@ -2761,7 +3419,10 @@ def sparse_attn_indexer( self.layer_idx) indexer_max_seq_len = metadata.get_indexer_max_seq_len() - if self.use_cute_dsl_paged_mqa_logits: + # A ragged batch cannot satisfy the DSL path's uniform kv_len + # broadcast / factor * atom == next_n; prepare_for_spec_decode + # routed it to the expanded buffers, keep the forward consistent. + if self.use_cute_dsl_paged_mqa_logits and not metadata.is_ragged_verify: # DSL kernel design: 1 atom per q (atom = real next_n positions), # kNumNextNAtoms = 1 for any real next_n. The matching schedule # is `scheduler_metadata_buffer` — built in `Indexer.prepare()` @@ -2880,7 +3541,28 @@ def sparse_attn_indexer( metadata.heuristic_scratch_values[ :num_gen_tokens] - if self.use_cute_dsl_topk and self._enable_heuristic_topk: + # The CuTe DSL top-k kernels reconstruct the row->request map + # from a scalar next_n and cannot express a ragged batch; + # route to the C++ kernel, which accepts the per-row extent + # directly. + ragged_row_kv_lens = metadata.ragged_row_kv_lens(num_gen_tokens) + if ragged_row_kv_lens is not None: + torch.ops.trtllm.indexer_topk_decode( + logits_decode, + gen_kv_lens_cuda, + topk_indices_buffer[num_ctx_tokens:num_ctx_tokens + + num_gen_tokens, :], + # next_n is dead on this path (the kernel reads + # row_kv_lens); pass 1 so a stale value cannot be used. + 1, + self.index_topk, + pre_idx=None, + heuristic_scratch=None, + compress_ratio=self.compress_ratio, + radix_aux_indices=metadata.radix_aux_indices, + radix_aux_logits=metadata.radix_aux_logits, + row_kv_lens=ragged_row_kv_lens) + elif self.use_cute_dsl_topk and self._enable_heuristic_topk: # GVR DSL: supports all compress_ratio and next_n values. torch.ops.trtllm.cute_dsl_gvr_topk_decode( logits_decode, diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 992eec88b7fd..7e41d33a111c 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -14,6 +14,7 @@ # limitations under the License. import functools +import contextlib import math import os import weakref @@ -517,6 +518,68 @@ def _bind_runtime_views( self.prompt_lens_cpu_runtime = prompt_lens_cpu self.host_request_types_runtime = host_request_types + def token_major_gen_view(self): + """The token-major generation presentation, or None. + + None on every backend but DSA, and on every uniform DSA batch. See + ``DSAtrtllmAttentionMetadata.token_major_gen_view``. + """ + return None + + @contextlib.contextmanager + def presented_token_major(self, view): + """Temporarily present the generation half as one row per query token. + + The MLA rope and sparse-MLA generation ops learn their sequence count + from ``host_context_lengths.size(0)`` and then require the query tokens + to divide evenly across it. A ragged batch does not divide. Swapping in + per-token views for the duration of those two calls makes the division + trivial (seq_len == 1) and satisfies their asserts rather than weakening + them. + + Scoped deliberately: ``num_seqs`` is read as a *request* count almost + everywhere else -- spec workers sizing acceptance rectangles, the + KV-length correction under the overlap scheduler, the indexer's own + metadata -- so the token-major rows must never outlive these two calls. + """ + if view is None: + yield + return + saved = (self.kv_lens_cuda_runtime, self.kv_lens_runtime, + self.prompt_lens_cpu_runtime, self.kv_cache_block_offsets, + self.max_num_requests, self.prompt_lens_cuda_runtime, + self.host_request_types_runtime) + try: + self.kv_lens_cuda_runtime = view.sequence_length + self.kv_lens_runtime = view.host_past_key_value_lengths + self.prompt_lens_cpu_runtime = view.host_context_lengths + self.kv_cache_block_offsets = view.kv_cache_block_offsets + # The op reserves its multi-CTA-KV counter as + # `num_heads * max_num_requests` and sizes its generation workspace + # from the same number. Under this presentation the batch dimension + # the kernels see is the ROW count, which is larger. Passing the + # static ceiling grows the reservation to cover it while keeping the + # op's (beam_width, max_num_requests, window) cache key constant. + # This is only a capacity input -- the op's own mMaxNumRequests is a + # JIT-warmup hint, and the per-step request count comes from the + # batch, not from here. + # + # The attention workspace is sized from the same number, so it + # grows with the row ceiling; that memory would otherwise hold + # KV cache. + self.max_num_requests = view.max_num_rows + # Every runtime view has to agree on batch_size with + # kv_lens_cuda_runtime, and the op reads request_types over all of + # num_seqs -- both are the row count here, not the request count. + self.prompt_lens_cuda_runtime = view.prompt_lens_cuda + self.host_request_types_runtime = view.host_request_types + yield + finally: + (self.kv_lens_cuda_runtime, self.kv_lens_runtime, + self.prompt_lens_cpu_runtime, self.kv_cache_block_offsets, + self.max_num_requests, self.prompt_lens_cuda_runtime, + self.host_request_types_runtime) = saved + def prepare(self) -> None: super().prepare() extra_attrs = get_model_extra_attrs() @@ -1770,13 +1833,21 @@ def forward( if not self.fmha_libs: self.create_fmha_libs() - for fmha in self.fmha_libs: - if fmha.is_supported(q, k, v, metadata, forward_args): - fmha.forward(q, k, v, metadata, forward_args) - break - else: - raise RuntimeError( - "No TRT-LLM attention FMHA library supports this request.") + # On a ragged generation step the attention op is handed + # one row per query token so its `num_tokens % num_seqs == 0` check + # passes with seq_len == 1. Wrapping the dispatch rather than the + # library keeps every fmha implementation, and its tests, untouched. + token_major_view = (metadata.token_major_gen_view() if + forward_args.attention_input_type + == AttentionInputType.generation_only else None) + with metadata.presented_token_major(token_major_view): + for fmha in self.fmha_libs: + if fmha.is_supported(q, k, v, metadata, forward_args): + fmha.forward(q, k, v, metadata, forward_args) + break + else: + raise RuntimeError( + "No TRT-LLM attention FMHA library supports this request.") if self.print_skip_softmax_stat: total_blocks, skipped_blocks = self.skip_softmax_stat @@ -2035,6 +2106,26 @@ def mla_rope_generation( metadata.helix_position_offsets, metadata.helix_is_inactive_rank ] + # Same presentation as the attention dispatch above: the + # rope kernel derives batch_idx = tok / seq_len and the KV write offset + # from a single scalar seq_len, which only describes a batch where every + # request contributes the same number of tokens. + with self.presented_token_major_for(metadata): + self._mla_rope_generation_impl( + metadata, fused_q, q_pe, latent_cache, cu_q_seqlens, + cu_kv_seqlens, fmha_scheduler_counter, mla_bmm1_scale, + mla_bmm2_scale, quant_q_buffer, out_scale, helix_tensor_params) + + @contextlib.contextmanager + def presented_token_major_for(self, metadata): + with metadata.presented_token_major(metadata.token_major_gen_view()): + yield + + def _mla_rope_generation_impl(self, metadata, fused_q, q_pe, latent_cache, + cu_q_seqlens, cu_kv_seqlens, + fmha_scheduler_counter, mla_bmm1_scale, + mla_bmm2_scale, quant_q_buffer, out_scale, + helix_tensor_params): torch.ops.trtllm.mla_rope_generation( fused_q, q_pe, diff --git a/tensorrt_llm/_torch/models/dspark/attention.py b/tensorrt_llm/_torch/models/dspark/attention.py index 7f63a40edfbe..0908ec90d05b 100644 --- a/tensorrt_llm/_torch/models/dspark/attention.py +++ b/tensorrt_llm/_torch/models/dspark/attention.py @@ -15,20 +15,13 @@ # # The DSpark captured-context attention primitives are ported from DeepSeek's # DeepSpec reference ``inference/kernel.py`` (``sparse_attn``) and -# ``inference/model.py`` (``get_dspark_topk_idxs``). The reference computes these -# with a TileLang kernel; this is a functional-first pure-PyTorch port with the -# same math (index-gather + online softmax + a learnable attention sink that -# contributes only to the softmax denominator). +# ``inference/model.py`` (``get_dspark_topk_idxs``); pure-PyTorch, same math. """DSpark draft captured-context attention primitives (hardware-agnostic). -The DSpark draft uses *dense* sliding-window MLA (``compress_ratio == 0``): the -query comes from the block's draft tokens, while the keys/values are gathered -from a small per-request set of positions (a sliding window of the projected -captured context plus the current block's own positions). Two primitives capture -the parts that differ from the standard MLA path: - -* :func:`get_dspark_topk_idxs` — the (window-context + block) position list. -* :func:`dspark_sparse_attn` — index-gathered attention with an attention sink. +Dense sliding-window MLA (``compress_ratio == 0``) for the draft block: queries +come from the block's draft tokens; keys/values are gathered from a rolling +window of projected captured context plus the block's own positions, with an +attention-sink softmax. """ from functools import lru_cache @@ -54,11 +47,8 @@ def precompute_dspark_freqs_cis( rope_theta: float = 10000.0, device: torch.device | str = "cpu", ) -> torch.Tensor: - """Plain (non-YaRN) RoPE complex exponentials for the DSpark draft. - - The dense draft attention (``compress_ratio == 0``) disables YaRN and uses the - base ``rope_theta`` (DeepSpec ``precompute_freqs_cis`` with - ``original_seq_len == 0``). + """Plain (non-YaRN) RoPE table: the dense draft disables YaRN and uses the + base ``rope_theta``. Returns: complex64 tensor ``[seqlen, rope_head_dim // 2]``. @@ -75,13 +65,11 @@ def precompute_dspark_freqs_cis( def apply_dspark_rotary( x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False ) -> torch.Tensor: - """Apply (or, with ``inverse``, de-apply) rotary embeddings, DeepSpec-style. + """Apply (or, with ``inverse``, de-apply) rotary embeddings. - Functional (non-in-place) port of DeepSpec ``apply_rotary_emb``: treats the - last dim as adjacent (re, im) pairs, rotates by ``freqs_cis`` indexed along the - sequence axis, and conjugates for the inverse (de-rotation applied to the - attention output). ``x`` is the rope-dim slice only: ``[b, s, rd]`` (3D) or - ``[b, s, h, rd]`` (4D), with ``freqs_cis`` of shape ``[s, rd // 2]``. + Last dim is adjacent (re, im) pairs; ``inverse`` conjugates. ``x`` is the + rope-dim slice only: ``[b, s, rd]`` or ``[b, s, h, rd]``, with ``freqs_cis`` + of shape ``[s, rd // 2]``. """ orig_dtype = x.dtype xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) @@ -98,13 +86,10 @@ def apply_dspark_rotary( def apply_dspark_rotary_batched( x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False ) -> torch.Tensor: - """Per-row (batched) variant of :func:`apply_dspark_rotary`. - - Identical math, but ``freqs_cis`` carries a leading batch axis so each row of - ``x`` is rotated by its own per-request phases (the generation draft runs each - request at a different absolute ``start_pos``). ``x`` is the rope-dim slice - only: ``[G, s, rd]`` (3D) or ``[G, s, h, rd]`` (4D), with ``freqs_cis`` of shape - ``[G, s, rd // 2]``. + """Per-row variant of :func:`apply_dspark_rotary`: ``freqs_cis`` carries a + leading batch axis (one set of phases per request). ``x`` is the rope-dim + slice only: ``[G, s, rd]`` or ``[G, s, h, rd]``, with ``freqs_cis`` of + shape ``[G, s, rd // 2]``. """ orig_dtype = x.dtype xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) @@ -138,19 +123,10 @@ def get_dspark_topk_idxs( ) -> torch.Tensor: """Per-query attended-position indices for the DSpark draft block. - Mirrors DeepSpec ``get_dspark_topk_idxs``: every one of the ``block_size`` - query positions attends to the same set — the ``min(window_size, start_pos+1)`` - most-recent context positions in the rolling KV window, then the - ``block_size`` positions of the current block (stored at offset - ``window_size`` in the concatenated KV). Note this is *non-causal* within the - block (every position sees every block position), matching the reference. - - Args: - window_size: sliding-window length of the captured-context KV cache. - bsz: batch size. - block_size: number of draft positions per request. - start_pos: absolute decode position (must be > 0); bounds the context. - device: device for the returned index tensor. + Every query position attends to the same set: the + ``min(window_size, start_pos+1)`` most-recent context positions in the + rolling KV window, then the ``block_size`` block positions stored at offset + ``window_size``. *Non-causal* within the block, matching the reference. Returns: int32 tensor ``[bsz, block_size, topk]`` with @@ -166,26 +142,13 @@ def get_dspark_topk_idxs_batched( block_size: int, start_pos: torch.Tensor, ) -> torch.Tensor: - """Sync-free, fixed-size (CUDA-graph-safe) batched ``get_dspark_topk_idxs``. - - Unlike the scalar :func:`get_dspark_topk_idxs` (whose ``topk`` width - ``min(window_size, start_pos+1) + block_size`` depends on the host int - ``start_pos``), this always returns the **fixed** width ``window_size + - block_size`` and masks the unfilled context slots with ``-1``. The masked - slots are excluded by :func:`dspark_sparse_attn` exactly as if they were - absent, so the result is numerically identical to gathering only the - ``min(window_size, start_pos+1)`` valid context positions — but the shape no - longer depends on the data, which is what CUDA-graph capture requires. - - Every query position attends to the same set: context window slots - ``0..window_size-1`` (slot ``c`` valid iff ``c <= start_pos[g]``, i.e. it has - been written) followed by the ``block_size`` block positions at offset - ``window_size`` (always valid). + """Sync-free, fixed-shape (CUDA-graph-safe) batched ``get_dspark_topk_idxs``. - Args: - window_size: sliding-window length of the captured-context KV cache. - block_size: number of draft positions per request. - start_pos: ``[G]`` int tensor of per-request absolute decode positions. + Always returns the fixed width ``window_size + block_size``, masking + unfilled context slots with ``-1`` (excluded by :func:`dspark_sparse_attn`), + so the result matches the scalar path while the shape stays data-independent + for graph capture. Context slot ``c`` is valid iff ``c <= start_pos[g]``; + ``start_pos`` is a ``[G]`` int tensor of absolute decode positions. Returns: int32 tensor ``[G, block_size, window_size + block_size]``. @@ -193,8 +156,6 @@ def get_dspark_topk_idxs_batched( device = start_pos.device g = start_pos.shape[0] ctx_cols = torch.arange(window_size, device=device) # [win] - # Context slot c holds a written key iff c <= start_pos (slots 0..start_pos - # filled; for start_pos >= window_size-1 the whole rolling window is filled). valid = ctx_cols.unsqueeze(0) <= start_pos.unsqueeze(1) # [G, win] ctx_idx = torch.where( valid, ctx_cols.unsqueeze(0).expand(g, -1), torch.full_like(valid, -1, dtype=torch.long) @@ -214,12 +175,10 @@ def dspark_sparse_attn( ) -> torch.Tensor: """Index-gathered multi-query attention with an attention sink. - Functional-first port of the DeepSpec ``sparse_attn`` TileLang kernel. For - each ``(batch, query, head)`` it gathers the ``topk`` KV rows named by - ``topk_idxs`` (an index of ``-1`` masks that slot), computes a scaled - dot-product softmax over them, and adds a per-head learnable *sink* logit that - participates only in the softmax denominator (i.e. an "attend-to-nothing" - option with a zero value vector). KV is shared across query heads (MQA). + Gathers the ``topk`` KV rows named by ``topk_idxs`` (``-1`` masks a slot), + softmaxes over them, and adds a per-head learnable *sink* logit that + participates only in the softmax denominator (an "attend-to-nothing" option + with a zero value vector). KV is shared across query heads (MQA). Args: q: ``[b, m, h, d]`` query (``m`` = block_size, ``h`` = heads). @@ -309,21 +268,15 @@ def dspark_attention_forward( ) -> torch.Tensor: """Captured-context DSpark draft attention (generation path, ``start_pos > 0``). - Functional port of DeepSpec ``DSparkAttention.forward`` for the dense - (``compress_ratio == 0``) draft: low-rank Q (``wq_a`` -> ``q_norm`` -> ``wq_b``) - with a per-head RMS + RoPE, MQA K/V from ``wkv`` (shared across heads), keys - gathered from a rolling captured-context window (``kv_cache``, into which the - projected ``main_x`` context is written at ``start_pos % window_size``) plus the - block's own positions, attention-sink softmax, inverse-RoPE on the output, and a - grouped low-rank O projection (``wo_a`` einsum + ``wo_b``). - - Weights are plain tensors for ``F.linear`` (the caller supplies the loaded / - dequantized projection weights); ``wo_a`` is the raw grouped weight matrix - ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` is - ``[b, window_size, head_dim]`` and is updated functionally (cloned). + Port of DeepSpec ``DSparkAttention.forward`` for the dense + (``compress_ratio == 0``) draft. Weights are plain (already dequantized) + tensors for ``F.linear``; ``wo_a`` is the raw grouped weight + ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` + is ``[b, window_size, head_dim]``; the projected ``main_x`` context is + written at ``start_pos % window_size`` (into a clone unless ``persist``). Returns: - ``[b, block_size, dim]`` attention output (residual stream contribution). + ``[b, block_size, dim]`` attention output. """ assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" b, block, _ = x.shape @@ -347,10 +300,8 @@ def dspark_attention_forward( kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [b, block, head_dim] kv = _rope_last_dims(kv, rd, blk_freqs) - # Write the context K/V into the rolling window, then attend over - # [window context | block] with the sink. ``persist=True`` writes through - # to the caller's buffer (cross-step decode, worker-owned window); the - # default clones so single-shot callers (golden / unit tests) stay pure. + # persist=True writes through to the caller's window (cross-step decode); + # the default clones so single-shot callers stay pure. cache = kv_cache if persist else kv_cache.clone() cache[:, start_pos % window_size] = main_kv.squeeze(1) kv_full = torch.cat([cache, kv], dim=1) # [b, window + block, head_dim] @@ -391,35 +342,25 @@ def dspark_attention_forward_batched( freqs_cis: torch.Tensor, persist: bool = False, ) -> torch.Tensor: - """Batched, CUDA-graph-safe captured-context DSpark draft attention. - - Numerically identical, per request, to :func:`dspark_attention_forward`, but - free of host syncs and data-dependent shapes so it can be captured into a CUDA - graph (the one-engine drafter runs inside the target's graph). The differences - from the scalar path are purely mechanical: - - * ``start_pos`` is a ``[G]`` int tensor (one absolute decode position per gen - request) instead of a python int; RoPE phases are *gathered* per request from - the fixed ``freqs_cis`` table rather than sliced. - * the rolling-window context K/V is written/read through the ``slots`` index - into a shared ``kv_cache`` (``persist=True`` writes through to the caller's - worker-owned buffer; otherwise a clone is used), instead of mutating a - per-request cache in place. - * the attended-position list has the fixed width ``window_size + block_size`` - with ``-1`` masking (see :func:`get_dspark_topk_idxs_batched`). + """Batched, CUDA-graph-safe :func:`dspark_attention_forward`. + + Numerically identical per request to the scalar path, but free of host + syncs and data-dependent shapes so it can be captured into a CUDA graph + (the one-engine drafter runs inside the target's graph). Args: x: ``[G, block, dim]`` block layer input (per gen request). main_x: ``[G, 1, hidden]`` projected captured context. start_pos: ``[G]`` int tensor of absolute decode positions (> 0). - kv_cache: ``[N, window_size, head_dim]`` rolling captured-context windows - (``N`` rows indexed by ``slots``; ``N == G`` for single-shot callers). + kv_cache: ``[N, window_size, head_dim]`` rolling captured-context + windows, rows indexed by ``slots`` (``persist=True`` writes + through to the caller's buffer; otherwise a clone is used). slots: ``[G]`` int tensor mapping each request to its ``kv_cache`` row. - freqs_cis: ``[maxlen, rope_head_dim // 2]`` precomputed plain-RoPE table; - must satisfy ``maxlen > start_pos.max() + block_size``. + freqs_cis: ``[maxlen, rope_head_dim // 2]`` plain-RoPE table; must + satisfy ``maxlen > start_pos.max() + block_size``. Returns: - ``[G, block, dim]`` attention output (residual stream contribution). + ``[G, block, dim]`` attention output. """ g, block, _ = x.shape rd = rope_head_dim @@ -442,10 +383,8 @@ def dspark_attention_forward_batched( kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [G, block, head_dim] kv = _rope_last_dims_batched(kv, rd, blk_freqs) - # Write the context K/V into the rolling window at slot start_pos%window_size, - # then attend over [window context | block]. ``persist=True`` writes through to - # the worker-owned buffer (cross-step decode); otherwise clone so single-shot - # callers stay pure. Indexed scatter/gather by (slots, slot_pos) is graph-safe. + # persist=True writes through to the worker-owned window; the default clones. + # Indexed scatter/gather by (slots, slot_pos) is graph-safe. write_target = kv_cache if persist else kv_cache.clone() slot_pos = start_pos % window_size # [G] write_target[slots, slot_pos] = main_kv.squeeze(1).to(write_target.dtype) diff --git a/tensorrt_llm/_torch/models/dspark/draft.py b/tensorrt_llm/_torch/models/dspark/draft.py index 1b47f4922965..ec02b64aeb83 100644 --- a/tensorrt_llm/_torch/models/dspark/draft.py +++ b/tensorrt_llm/_torch/models/dspark/draft.py @@ -24,7 +24,7 @@ - ``build_draft_input_ids``: ``[bonus_token, noise, noise, ...]`` block input. - ``dspark_propose``: given the per-position backbone ``base_logits`` and the Markov / confidence heads, run the autoregressive Markov refinement to sample - the block tokens and apply the static confidence-threshold truncation. + the block tokens and score them with the confidence head. The backbone (3 V4 blocks producing ``block_hidden``) lives in the model module; this file is the part fully specified by the reference and validated against it. @@ -35,8 +35,6 @@ import torch from torch import nn -from .heads import confident_prefix_length - def build_draft_input_ids( bonus_token_ids: torch.Tensor, *, block_size: int, noise_token_id: int @@ -61,29 +59,36 @@ def dspark_propose( confidence_head: Optional[nn.Module], block_size: int, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, ) -> tuple: """Produce DSpark draft tokens for one block (functional-first, static length). + Always proposes the *full* block: the confidence head does not shorten the + draft (the block is produced by a single parallel backbone pass, so there is + nothing to save there). It scores the block so the verification scheduler can + decide how many of those tokens are worth sending to the target. + Args: base_logits: ``[batch, block_size, vocab]`` from the backbone + lm_head. bonus_token_ids: ``[batch]`` the token preceding the first draft position. block_hidden: ``[batch, block_size, hidden]`` backbone hidden (feeds the confidence head, and the RNN-head variant). markov_head / confidence_head: the validated DSpark heads (may be None). + return_confidence: compute the per-position confidence logits. This is a + run-constant flag (read once from the decoding config), never a + per-step decision, so it cannot make the captured graph diverge. Returns: draft_tokens: ``[batch, block_size]`` sampled tokens (full block; callers keep the tensor fixed-width for CUDA-graph safety). - num_proposed: ``[batch]`` int32 — how many leading tokens survive the - static confidence-threshold truncation (== block_size when no head / - threshold<=0). + confidence: ``[batch, block_size]`` fp32 *raw* confidence logits, or None + when disabled / no head. Calibrate with ``confidence_head.apply_sts`` + before taking the cumulative product. """ - batch = base_logits.shape[0] # ``draft_logits`` are the per-position distributions the draft token is drawn # from (markov-corrected when a head is present, else the raw base logits). - # Surfaced under ``return_logits`` for the §7.9 probabilistic-acceptance - # (1-TV) measurement; the normal path ignores them. + # Surfaced under ``return_logits`` to feed the rejection-sampling verifier; + # the normal path ignores them. draft_logits = base_logits if markov_head is not None: draft_tokens, corrected = markov_head.sample_block_tokens( @@ -98,33 +103,22 @@ def dspark_propose( draft_tokens = greedy_or_sample(base_logits, temperature) - # Scaffolding: confidence-based dynamic drafting is NOT enabled in this PR. - # The worker always calls with confidence_threshold=0.0, so the block below is - # inert and num_proposed stays == block_size (the full block is proposed). The - # returned num_proposed is intentionally not yet consumed by the speculative - # scheduler/verifier; wiring it through is a follow-up (see PR description). - num_proposed = torch.full( - (batch,), int(block_size), dtype=torch.int32, device=base_logits.device - ) - if confidence_head is not None and confidence_threshold > 0.0: + # Confidence scoring: branch-free and fixed-shape (no ``.item()``, no + # data-dependent shapes), so the whole block stays CUDA-graph capturable. + confidence = None + if return_confidence and confidence_head is not None: # prev token at position k is [bonus, draft_0, ..., draft_{k-1}] prev_ids = torch.cat([bonus_token_ids.unsqueeze(1), draft_tokens[:, :-1]], dim=1) prev_emb = ( markov_head.get_prev_embeddings(prev_ids) - if (markov_head is not None and getattr(confidence_head, "with_markov", False)) + if (markov_head is not None and confidence_head.with_markov) else None ) - conf_logits = ( + confidence = ( confidence_head(block_hidden, prev_embeddings=prev_emb) if prev_emb is not None else confidence_head(block_hidden) ) - # Per-request prefix truncation (batch handled row-wise to stay simple; - # functional-first scope typically runs batch=1 for the draft). - for b in range(batch): - num_proposed[b] = confident_prefix_length( - conf_logits[b : b + 1], block_size=block_size, threshold=confidence_threshold - ) if return_logits: - return draft_tokens, num_proposed, draft_logits - return draft_tokens, num_proposed + return draft_tokens, confidence, draft_logits + return draft_tokens, confidence diff --git a/tensorrt_llm/_torch/models/dspark/heads.py b/tensorrt_llm/_torch/models/dspark/heads.py index c49e35fbafa0..e2010c0f1436 100644 --- a/tensorrt_llm/_torch/models/dspark/heads.py +++ b/tensorrt_llm/_torch/models/dspark/heads.py @@ -15,21 +15,13 @@ # # The DSpark Markov/RNN/confidence-head math is ported from DeepSeek's DeepSpec # reference implementation (https://github.com/deepseek-ai/DeepSpec, MIT License). -"""DSpark draft-network heads (pure-torch, framework-agnostic). +"""DSpark draft-network heads (pure-torch so they unit-test against DeepSpec). -These modules implement the *sequential refinement* and *acceptance-confidence* -parts of DeepSeek's DSpark speculative-decoding draft network: - - - Markov head: a low-rank token-bigram logit bias ``logits_k += W2(W1[t_{k-1}])`` - applied autoregressively across the ``block_size`` draft positions (the cheap - "sequential" half of DSpark's "semi-parallel" drafting). RNN variant carries - a GRU-style recurrent state across positions. - - Confidence head: predicts a per-position acceptance probability; the cumulative - product over positions estimates prefix-acceptance and is used only to - *truncate* the proposed draft length (NOT to decide acceptance). - -This file deliberately depends on ``torch`` only so it can be unit-tested in -isolation (token-for-token) against the DeepSpec reference. +Markov/RNN heads apply an autoregressive per-position logit bias across the +draft block. The confidence head predicts per-position *conditional* acceptance +probabilities (``P(accept_k | accept_1..k-1)``) whose cumulative product the +verification scheduler budgets against; it decides *how many* tokens are +verified, never *whether* one is accepted, so scheduling stays lossless. """ from typing import Optional @@ -213,22 +205,41 @@ def build_markov_head( class DSparkConfidenceHead(nn.Module): """Per-position acceptance-confidence predictor (DeepSpec AcceptRatePredictor). - Input features are the backbone hidden state, optionally concatenated with the - Markov head's previous-token embedding. Output is a single logit per position. + Emits one raw logit per position; :meth:`apply_sts` turns it into a + probability, folding in the sequential-temperature-scaling (STS) + calibration (``sts_temperatures`` defaults to all-ones: plain sigmoid). """ - def __init__(self, *, hidden_size: int, markov_rank: int = 0, with_markov: bool = False): + def __init__( + self, + *, + hidden_size: int, + markov_rank: int = 0, + with_markov: bool = False, + bias: bool = False, + block_size: int = 0, + ): super().__init__() self.with_markov = bool(with_markov) input_dim = int(hidden_size) + (int(markov_rank) if with_markov else 0) - # The checkpoint stores ``proj`` as a bias-free bf16 weight, but the - # confidence score is computed in fp32 (mirrors the DeepSpec reference - # ``Linear(input_dim, 1, dtype=torch.float32)`` with the fp32 matmul). - self.proj = nn.Linear(input_dim, 1, bias=False, dtype=torch.float32) + # Checkpoint weight is bf16, but the score matmul is fp32 (DeepSpec parity). + self.proj = nn.Linear(input_dim, 1, bias=bool(bias), dtype=torch.float32) + # Per-position temperatures, broadcast against ``[batch, block]``. + # Update in place only: a captured CUDA graph holds this storage; + # rebinding the attribute would be invisible to the graph. + self.register_buffer( + "sts_temperatures", + torch.ones(max(int(block_size), 1), dtype=torch.float32), + persistent=False, + ) + # CPU mirror of ``sts_temperatures``; rebuilt lazily by + # ``_host_sts_temperatures`` and dropped by ``load_sts_temperatures``. + self._sts_temperatures_host: Optional[torch.Tensor] = None def forward( self, hidden_states: torch.Tensor, prev_embeddings: Optional[torch.Tensor] = None ) -> torch.Tensor: + """``[*, block, hidden] -> [*, block]`` raw (uncalibrated) logits.""" if self.with_markov: assert prev_embeddings is not None features = torch.cat([hidden_states, prev_embeddings.to(hidden_states.dtype)], dim=-1) @@ -237,18 +248,57 @@ def forward( # fp32 matmul for a stable confidence score (mirrors the reference). return self.proj(features.float()).squeeze(-1) + def apply_sts(self, confidence_logits: torch.Tensor) -> torch.Tensor: + """Raw logits -> calibrated per-position acceptance probabilities in [0, 1]. -def confident_prefix_length( - confidence_logits: torch.Tensor, *, block_size: int, threshold: float -) -> int: - """First position k where ``sigmoid(confidence_k) < threshold``. - - Returns ``block_size`` when threshold<=0 (no truncation) or all positions - are confident. Assumes batch size 1 (functional-first scope). - """ - if threshold <= 0.0: - return int(block_size) - below = confidence_logits.sigmoid() < threshold - if not bool(below[0].any().item()): - return int(block_size) - return int(torch.nonzero(below[0], as_tuple=False)[0].item()) + Must accept CPU logits: the verification planner calibrates on pinned + host memory while the head lives on the device. The device-mismatch + branch is host-side and devices match on the in-graph path, so nothing + extra is captured. + """ + temperatures = self.sts_temperatures + if temperatures.device != confidence_logits.device: + if confidence_logits.device.type == "cpu": + temperatures = self._host_sts_temperatures() + else: + temperatures = temperatures.to(confidence_logits.device) + return torch.sigmoid(confidence_logits.float() / temperatures) + + def _host_sts_temperatures(self) -> torch.Tensor: + """CPU mirror of ``sts_temperatures``, materialized at most once.""" + cached = self._sts_temperatures_host + if cached is None: + cached = self.sts_temperatures.detach().to("cpu") + self._sts_temperatures_host = cached + return cached + + @torch.no_grad() + def load_sts_temperatures(self, temperatures: torch.Tensor) -> None: + """In-place update of the STS table (never rebind: see ``__init__``).""" + flat = temperatures.reshape(-1).to(device=self.sts_temperatures.device, dtype=torch.float32) + if flat.numel() != self.sts_temperatures.numel(): + raise ValueError( + f"STS temperature table has {flat.numel()} entries but the confidence " + f"head expects {self.sts_temperatures.numel()} (one per block position)" + ) + if not bool(torch.all(flat > 0.0)): + raise ValueError("STS temperatures must be strictly positive") + self.sts_temperatures.copy_(flat) + # Drop the CPU mirror; apply_sts rebuilds it on next use. + self._sts_temperatures_host = None + + def load_weights(self, weights: list) -> None: + """Strict loader: fail loudly on checkpoint keys the module lacks (a + silently dropped ``proj.bias`` would shift every confidence score).""" + (module_weights,) = weights + expected = dict(self.named_parameters()) + unexpected = [k for k in module_weights if k not in expected] + if unexpected: + raise ValueError( + f"DSpark confidence head checkpoint has unsupported key(s) {sorted(unexpected)}; " + f"the module exposes {sorted(expected)}. A checkpoint bias requires constructing " + f"DSparkConfidenceHead(bias=True) -- silently dropping it would bias every " + f"per-position confidence score." + ) + for name, param in expected.items(): + param.data.copy_(module_weights[name][:]) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 4efadd01784d..b46ef07aeffa 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -146,15 +146,6 @@ def remap_dspark_draft_keys(weights: Dict, num_stages: int) -> Dict: return out -# The checkpoint stores -# ``mtp.{s}.attn.wo_a`` as fp8_e4m3 + a UE8M0 128x128 block scale (verified), and -# the reference (`inference/model.py`, ``self.wo_a`` is a bf16 ColumnParallelLinear -# loaded from the fp8 ckpt) uses the DEQUANTIZED bf16 ``wo_a`` (== ``wo_a_fp8 * -# scale`` ~ absmean 0.065). The bf16 captured-context path historically skipped -# this dequant (raw fp8-cast-to-bf16, ~993x too large); the correct behavior is to -# dequantize ``wo_a`` (cos 1.0 vs ``wo_a_fp8 * scale``). Always dequantize now. - - class DSparkBlock(DeepseekV4DecoderLayer): """One DSpark draft stage = a DeepSeek-V4 decoder block + DSpark extras. @@ -171,6 +162,7 @@ def __init__( stage_id: int, num_stages: int, num_capture_layers: int, + block_size: int = 0, ): # The inherited attention uses a draft-local layer index, while the # decoder layer keeps its model-level index for weights and captures. @@ -182,7 +174,7 @@ def __init__( disable_post_moe_fusion=True, ) config = model_config.pretrained_config - spec_cfg = getattr(model_config, "spec_config", None) + spec_cfg = model_config.spec_config self.stage_id = int(stage_id) self.num_stages = int(num_stages) # mask_token_id is a user override on the speculative_config; None means @@ -236,6 +228,8 @@ def __init__( # markov_rank <= 0); otherwise dspark_propose passes no # prev_embeddings and DSparkConfidenceHead.forward would assert. with_markov=self.markov_rank > 0, + # Sizes the per-position STS calibration table once, up front. + block_size=int(block_size), ) @property @@ -272,7 +266,7 @@ def __init__( # ``mtp.{0..n-1}.*`` weight namespace. Resolve it from (in priority): # an explicit override, the spec config's ``num_draft_layers``, a # pretrained-config ``n_mtp_layers``, else fall back to nextn. - spec_cfg = getattr(model_config, "spec_config", None) + spec_cfg = model_config.spec_config self.num_stages = int( num_stages if num_stages is not None @@ -326,6 +320,7 @@ def __init__( stage_id=s, num_stages=self.num_stages, num_capture_layers=self.num_capture_layers, + block_size=self.block_size, ) for s in range(self.num_stages) ] @@ -437,7 +432,7 @@ def cache_attn_weights_from_checkpoint(self, ckpt_dir: str, weight_map: Dict[str """Populate ``_dspark_attn`` by reading the ``mtp.{s}.attn.*`` tensors from the checkpoint shards on disk, then dequantizing via :meth:`_cache_attn_weights`. - TODO(step 3): source these from the loaded ``MLA`` modules instead, once the + TODO: source these from the loaded ``MLA`` modules instead, once the fused/interleaved fp8 scale layout is decoded, to drop the checkpoint I/O. """ from safetensors import safe_open @@ -526,8 +521,8 @@ def _draft_normalized_quant_config(model_config): when nothing changes (e.g. the MXFP4 checkpoint, whose global algo is not ``MIXED_PRECISION``) so the draft config is left byte-identical. """ - qc = getattr(model_config, "quant_config", None) - if qc is None or getattr(qc, "quant_algo", None) != QuantAlgo.MIXED_PRECISION: + qc = model_config.quant_config + if qc is None or qc.quant_algo != QuantAlgo.MIXED_PRECISION: return None # Reuse the target normalizer on a throwaway shallow copy so we only # extract the resolved global quant_config; the shared ``model_config`` @@ -535,7 +530,7 @@ def _draft_normalized_quant_config(model_config): probe = copy.copy(model_config) object.__setattr__(probe, "_frozen", False) normalized = _normalize_deepseek_v4_nvfp4_mixed_precision_config(probe) - resolved = getattr(normalized, "quant_config", qc) + resolved = normalized.quant_config return resolved if resolved is not qc else None @staticmethod @@ -546,7 +541,7 @@ def _draft_sparse_config(model_config, base: int, num_stages: int): sparse config must expose the draft layers' per-layer ratios at indices ``[0, num_stages)``. """ - sa = getattr(model_config, "sparse_attention_config", None) + sa = model_config.sparse_attention_config compress_ratios = getattr(sa, "compress_ratios", None) if sa is not None else None if not compress_ratios or len(compress_ratios) < base + num_stages: return None @@ -569,7 +564,7 @@ def _draft_quant_config_dict(model_config, base: int, num_stages: int): physically MXFP4 (identical to the main MoE layers), so copy a representative main MoE layer's experts quant onto the draft layer keys. """ - qcd = getattr(model_config, "quant_config_dict", None) + qcd = model_config.quant_config_dict if not qcd: return None src = next( @@ -829,10 +824,7 @@ def _forward_stage( # distinct set of requests) and needs no cross-rank reduction; but under # plain tensor parallelism (attention_dp off, tp_size > 1) the expert-sharded # MoE output must be all-reduced across ranks -- exactly what the target MoE - # does (enable_allreduce = not (POST_MOE_FUSION or tp_size == 1)). Previously - # this was hard-coded to False, which dropped the reduction on the non-ADP - # path, corrupting the draft block proposals and roughly halving DSpark - # acceptance length whenever attention_dp was disabled. + # does (enable_allreduce = not (POST_MOE_FUSION or tp_size == 1)). moe_enable_allreduce = ( not self.model_config.mapping.enable_attention_dp and self.model_config.mapping.tp_size > 1 @@ -857,7 +849,7 @@ def forward( *, kv_windows: Optional[torch.Tensor] = None, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, all_rank_num_tokens: Optional[List[int]] = None, ) -> tuple: @@ -877,7 +869,7 @@ def forward( worker; updated in place each call. ``None`` allocates fresh zero windows (single-shot golden / test path). Returns: - ``(draft_tokens [T, block], num_proposed [T])`` from ``forward_head``. + ``(draft_tokens [T, block], confidence [T, block] or None)``. """ assert start_pos > 0, "DSpark draft runs at generation (start_pos > 0)" if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: @@ -908,7 +900,7 @@ def forward( h, bonus_token_ids, temperature=temperature, - confidence_threshold=confidence_threshold, + return_confidence=return_confidence, return_logits=return_logits, ) @@ -921,7 +913,7 @@ def forward_batched( kv_windows: torch.Tensor, slots: torch.Tensor, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, all_rank_num_tokens: Optional[List[int]] = None, ) -> tuple: @@ -933,9 +925,9 @@ def forward_batched( graph). ``start_pos`` is a ``[G]`` tensor (one absolute decode position per gen request); the rolling captured-context windows are written/read through ``slots`` into the worker-owned ``kv_windows`` buffer; RoPE phases are - gathered from a fixed table. ``forward_head`` is run with - ``confidence_threshold == 0`` (the worker proposes the full block), which is - the graph-safe branch of :func:`dspark_propose`. + gathered from a fixed table. ``forward_head`` always proposes the full + block; ``return_confidence`` only adds a fixed-shape scoring matmul, so + both settings stay capturable. Args: main_hidden: ``[G, num_capture * hidden]`` captured target context. @@ -945,7 +937,7 @@ def forward_batched( windows; written in place through ``slots``. slots: ``[G]`` int tensor mapping each request to its ``kv_windows`` row. Returns: - ``(draft_tokens [G, block], num_proposed [G])`` from ``forward_head``. + ``(draft_tokens [G, block], confidence [G, block] or None)``. """ if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: raise RuntimeError( @@ -976,7 +968,7 @@ def forward_batched( h, bonus_token_ids, temperature=temperature, - confidence_threshold=confidence_threshold, + return_confidence=return_confidence, return_logits=return_logits, ) @@ -1025,14 +1017,14 @@ def forward_head( bonus_token_ids: torch.Tensor, *, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, ) -> tuple: """Block-draft head: hc_head + norm + lm_head -> markov refine + confidence. ``block_hidden`` is the last stage's mHC residual ``[*, block, hc_mult, hidden]``. - Returns (draft_tokens [*, block], num_proposed [*]); with ``return_logits`` - also returns the per-position draft logits [*, block, vocab] (§7.9 1-TV). + Returns (draft_tokens [*, block], confidence [*, block] or None); with ``return_logits`` + also returns the per-position draft logits [*, block, vocab]. """ last = self.mtp_layers[-1] h = last.hc_head(block_hidden) @@ -1046,7 +1038,7 @@ def forward_head( confidence_head=last.confidence_head, block_size=self.block_size, temperature=temperature, - confidence_threshold=confidence_threshold, + return_confidence=return_confidence, return_logits=return_logits, ) diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index 63f9a1bd65d3..639cd98ad81d 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -2523,8 +2523,17 @@ def forward_absorption_generation( # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp num_seqs = attn_metadata.num_seqs - cu_q_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=q.device) - cu_kv_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=q.device) + # The rope kernel writes seqQOffset[batch_idx + 1] for every + # batch_idx < batch_size, and CUB-scans total_s_len / seq_len entries + # into seqKVOffsets (mlaKernels.cu:579, :650). `batch_size` there is the + # count of presented rows, which equals num_seqs only while every + # request contributes the same number of query tokens. Under the + # token-major presentation a row is a token, so size these by the token + # count -- which is also an upper bound on num_seqs on every other path, + # so this is safe rather than merely sufficient. + cu_seqlens_rows = max(num_seqs, num_tokens) + 1 + cu_q_seqlens = torch.empty(cu_seqlens_rows, dtype=torch.int32, device=q.device) + cu_kv_seqlens = torch.empty(cu_seqlens_rows, dtype=torch.int32, device=q.device) fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=q.device) has_fp8_kv_cache = ( self.mqa.has_fp8_kv_cache if hasattr(self.mqa, "has_fp8_kv_cache") else False diff --git a/tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py b/tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py index dff8d078eedb..f0f011da724e 100644 --- a/tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py +++ b/tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py @@ -51,6 +51,11 @@ class ADPIterStatsRecord: host_step_time_ms: Optional[float] = None prev_device_step_time_ms: Optional[float] = None gpu_forward_time_ms: Optional[float] = None + #: Executed (padded) generation-token total, captured on rank 0 and + #: broadcast to every rank row like the timings above. Group-uniform by + #: construction on graph steps: every ADP rank pads and fills to the + #: group's agreed shape, so rank 0's value IS the group's. + num_generation_tokens: Optional[int] = None class ADPIterStatsBuffer: @@ -80,6 +85,7 @@ def __init__(self) -> None: self._rank0_host_step_time_ms: Dict[int, Optional[float]] = {} self._rank0_prev_device_step_time_ms: Dict[int, Optional[float]] = {} self._rank0_gpu_forward_time_ms: Dict[int, Optional[float]] = {} + self._rank0_num_generation_tokens: Dict[int, Optional[int]] = {} self._oldest_iter: Optional[int] = None @staticmethod @@ -108,6 +114,7 @@ def queue( host_step_time_ms: Optional[float] = None, prev_device_step_time_ms: Optional[float] = None, gpu_forward_time_ms: Optional[float] = None, + num_generation_tokens: Optional[int] = None, ) -> None: """Queue local stats; rank 0 also keeps objects needed for fanout.""" payload = self.make_payload(stats) @@ -129,6 +136,7 @@ def queue( self._rank0_host_step_time_ms[iter_id] = host_step_time_ms self._rank0_prev_device_step_time_ms[iter_id] = prev_device_step_time_ms self._rank0_gpu_forward_time_ms[iter_id] = gpu_forward_time_ms + self._rank0_num_generation_tokens[iter_id] = num_generation_tokens def next_payload(self) -> Optional[RankIterStatsPayload]: """Return the oldest pending stats payload to piggyback.""" @@ -163,6 +171,7 @@ def _discard(self, iter_id: int, *, recompute_oldest: bool = True) -> None: self._rank0_host_step_time_ms.pop(iter_id, None) self._rank0_prev_device_step_time_ms.pop(iter_id, None) self._rank0_gpu_forward_time_ms.pop(iter_id, None) + self._rank0_num_generation_tokens.pop(iter_id, None) if recompute_oldest and iter_id == self._oldest_iter: self._recompute_oldest_iter() @@ -294,6 +303,7 @@ def finalize( host_step_time_ms = self._rank0_host_step_time_ms.get(iter_stats_iter) prev_device_step_time_ms = self._rank0_prev_device_step_time_ms.get(iter_stats_iter) gpu_forward_time_ms = self._rank0_gpu_forward_time_ms.get(iter_stats_iter) + num_generation_tokens = self._rank0_num_generation_tokens.get(iter_stats_iter) for rank_state in sorted(matching_states, key=lambda s: s.rank): rank = rank_state.rank @@ -306,6 +316,7 @@ def finalize( host_step_time_ms=host_step_time_ms, prev_device_step_time_ms=prev_device_step_time_ms, gpu_forward_time_ms=gpu_forward_time_ms, + num_generation_tokens=num_generation_tokens, ) ) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index c2a6a5bb4de1..998c59b9c6d5 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1,3 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 import bisect import contextlib from dataclasses import dataclass @@ -135,6 +137,13 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.spec_config = config.spec_config self.sparse_config = config.sparse_attention_config self.is_encoder_decoder = config.is_encoder_decoder + # Window each ragged padding row carries, in py_verify_len units. Set + # by ModelEngine.fit_ragged_verify_lens when it picks the token bucket; + # the padding rows are one shared dummy object so they all share it. + self.ragged_pad_verify_len = 0 + # Published by ModelEngine when the ragged fit picks a token bucket; + # None is the resting value (no ragged step in flight). + self.agreed_ragged_bucket: Optional[int] = None self.graphs: Dict[KeyType, torch.cuda.CUDAGraph] = {} self.graph_outputs: Dict[KeyType, @@ -265,8 +274,52 @@ def get_graph_key( set(draft_len_list)) == 1, "All draft lengths must be the same" key = (batch_size, draft_len, False, short_seq_len_mode, is_all_greedy_sample) + # Ragged verification bakes the flat verified-token total into the + # captured graph, so it becomes a second key axis; the element is + # absent for uniform batches, whose keys stay exactly what they + # were. + verify_bucket = self._ragged_verify_bucket(batch) + if verify_bucket is not None: + key = key + (verify_bucket, ) return key + def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: + """Total verified token count when the batch uses ragged windows. + + None unless *every* generation request carries a window; None keeps + the original (uniform) key. Config-gated before touching the batch: + this runs on every graph-key build for every model. + """ + # Keyed on the MODE: `cap-accept` shares this config flag but never + # shrinks the token axis, so it must keep the ordinary uniform key. + from ..speculative.dspark_observability import trims_submitted_tokens + if not trims_submitted_tokens(self.spec_config): + return None + # Windows can vanish between the fit and key derivation; a ragged key + # over a window-less batch replays stale token-major row maps (IMA). + if any( + getattr(request, "py_verify_len", None) is None + for request in batch.generation_requests): + return None + # Return the fitted bucket, never a fresh sum over the batch: every + # attention-DP rank must key on the same value, and the fit's + # allgathered arithmetic agrees by construction where a local sum only + # agrees by accident. + return self.agreed_ragged_bucket + + @staticmethod + def _local_draft_len(batch: ScheduledRequests) -> int: + """This rank's draft length, for the attention-DP shape comparison. + + Tolerant where :meth:`get_graph_key` asserts: a malformed batch + reports its max, differs from well-formed peers, and the gate + declines the graph. + """ + return max( + (len(request.py_draft_tokens) + for request in batch.generation_requests), + default=0) + @staticmethod def _get_mrope_position_delta(request: Any) -> Optional[Any]: mrope_position_delta = getattr(request, "py_mrope_position_delta", None) @@ -314,25 +367,46 @@ def maybe_get_cuda_graph( - The spec_metadata for the graph, if applicable. - The key for the graph, if applicable. """ + # Why the last call declined, for the DSpark ragged counters; the call + # site that counts the miss cannot see the reason. + self.last_miss_reason = None + self.last_miss_key = None + # disable when doing statistic if ExpertStatistic.should_record(): + self.last_miss_reason = "expert_statistic" return None, None, None can_run_cuda_graph = batch.can_run_cuda_graph batch_size = batch.batch_size if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - all_can_graph_batch = self.config.dist.tp_allgather( - [can_run_cuda_graph, batch_size]) + # Compare every key component a rank can decide for itself: replay + # freezes all_rank_num_tokens, so a shape mismatch is invisible at + # replay time. draft_len is settable rank-locally (acceptance-rate + # gate), so it is compared rather than assumed. -1 = "not + # applicable" so the payload is plain ints on every rank. + local_bucket = self._ragged_verify_bucket( + batch) if can_run_cuda_graph else None + all_can_graph_batch = self.config.dist.tp_allgather([ + can_run_cuda_graph, batch_size, + -1 if local_bucket is None else int(local_bucket), + self._local_draft_len(batch) if can_run_cuda_graph else -1 + ]) is_all_gen_only = all(all_can_graph[0] for all_can_graph in all_can_graph_batch) - all_batch_size_equal = all( - all_gen_only[1] == all_can_graph_batch[0][1] - for all_gen_only in all_can_graph_batch) - - if not is_all_gen_only or not all_batch_size_equal: + all_shapes_equal = all( + peer[1:] == all_can_graph_batch[0][1:] + for peer in all_can_graph_batch) + + if not is_all_gen_only or not all_shapes_equal: + self.last_miss_reason = ("peer_not_gen_only" + if not is_all_gen_only else + "peer_shape_mismatch") return None, None, None if not self.enabled or not can_run_cuda_graph: + self.last_miss_reason = ("disabled" if not self.enabled else + "local_not_gen_only") return None, None, None if self.config.use_mrope and any( self._needs_mrope_delta_cache_update(request) @@ -341,6 +415,7 @@ def maybe_get_cuda_graph( # Qwen3.5 configs normalized to text-only decoding). Requests that # do carry a delta must first populate the model-side cache for their # current seq slot before graph replay. + self.last_miss_reason = "mrope_delta_cache" return None, None, None key = self.get_graph_key(batch, new_tensors_device, spec_resource_manager, spec_metadata) @@ -352,9 +427,13 @@ def maybe_get_cuda_graph( # Graph doesn't exist yet. If on-the-fly capture is not allowed, # fall back to eager so the caller doesn't need a separate check. if not self._capture_allowed: + self.last_miss_reason = "key_not_captured" + self.last_miss_key = key return None, None, None if batch_size not in self.supported_batch_sizes: + self.last_miss_reason = "bs_not_supported" + self.last_miss_key = key return None, None, None num_sequences_in_batch = batch_size * self.max_beam_width @@ -410,7 +489,10 @@ def capture(self, # We're forced to do this because we cannot reallocate inputs over many graph runs. max_draft_len = key[1] token_per_request = max_draft_len + 1 - num_tokens_for_capture = (batch_size * self.max_beam_width * + # Ragged keys carry the flat token total as a sixth element; that is + # the replay width, NOT bs * (draft_len + 1). + num_tokens_for_capture = (key[5] if len(key) > 5 else + batch_size * self.max_beam_width * token_per_request) sliced_static_tensors = { @@ -516,6 +598,25 @@ def replay(self, key: KeyType, return output_ref + def will_pad_to(self, padded_bs: int, num_requests: int) -> bool: + """Whether a batch of ``num_requests`` can actually reach ``padded_bs``. + + Mirrors the SIZE guards in :meth:`_get_padded_batch` only; the + resource-dependent bails (dummy KV capacity, encoder-decoder state) + are deliberately not predicted. The ragged fit sizes its bucket grid + from the padded row count, so a wrong yes strands the fitted total in + no captured bucket. + """ + if padded_bs <= 0 or num_requests <= 0: + return False + if padded_bs not in self.supported_batch_sizes: + return False + if padded_bs == num_requests: + return True # already at a captured size, nothing to add + # _get_padded_batch bails when padding_size + batch_size exceeds the + # configured batch size; that sum is just padded_bs. + return padded_bs <= self.config.batch_size + def _get_padded_batch(self, batch: ScheduledRequests, resource_manager: ResourceManager, runtime_draft_len: int) -> int: @@ -615,6 +716,26 @@ def _get_padded_batch(self, batch: ScheduledRequests, self.padding_dummy_requests[runtime_draft_len] = dummy_request padding_dummy_request = self.padding_dummy_requests[runtime_draft_len] + # The pad rows' (uniform) window was decided by the fit when it chose + # the bucket (they are one shared object), so read it back. Cleared + # otherwise: the dummy is cached across steps and would carry a stale + # window into a uniform batch. + if getattr(self.spec_config, "enable_ragged_verify", False): + ragged = any( + getattr(request, "py_verify_len", None) is not None + for request in batch.generation_requests) + padding_dummy_request.py_verify_len = (self.ragged_pad_verify_len + if ragged else None) + # Same for cap-accept: `_attach_accept_caps` drops the whole + # batch's caps if any request lacks one, so a capless dummy would + # silently degrade the mode to `static`. Its acceptance is + # discarded, so give it the full block; cleared otherwise against + # staleness. + cap_accept = any( + getattr(request, "py_verify_cap", None) is not None + for request in batch.generation_requests) + padding_dummy_request.py_verify_cap = ( + self.spec_config.max_draft_len if cap_accept else None) batch.generation_requests.extend([padding_dummy_request] * padding_size) return padding_size diff --git a/tensorrt_llm/_torch/pyexecutor/hang_detector.py b/tensorrt_llm/_torch/pyexecutor/hang_detector.py index 2ae692ed5902..f6dde7c58959 100644 --- a/tensorrt_llm/_torch/pyexecutor/hang_detector.py +++ b/tensorrt_llm/_torch/pyexecutor/hang_detector.py @@ -100,6 +100,7 @@ def __init__( self.lock = threading.Lock() self.active = False self._detected = False + self._status_providers: list[Callable[[], str]] = [] def start(self): """Enable hang detection.""" @@ -113,13 +114,38 @@ def run_loop(): self.loop_thread = threading.Thread(target=run_loop, daemon=True, name="hang_detector_loop") self.loop_thread.start() - async def _detect_hang(self): + def register_status_provider(self, provider: Callable[[], str]) -> None: + """Register a nonblocking callable that returns status to dump on hang detection.""" + with self.lock: + self._status_providers.append(provider) + + async def _detect_hang(self) -> None: await asyncio.sleep(self.timeout) with self.lock: - self._detected = True - logger.error(f"Hang detected after {self.timeout} seconds.") + status_providers = tuple(self._status_providers) + + # All diagnostics are best-effort: nothing may prevent on_detected() + # (hard-kill propagation) from firing. + _best_effort_log_error(f"Hang detected after {self.timeout} seconds.") + for provider in status_providers: + try: + status = provider() + if status: + _best_effort_log_error(status) + except Exception as error: # noqa: BLE001 - isolate diagnostic providers + _best_effort_log_error( + f"HangDetector: status provider failed with {type(error).__name__}: {error}" + ) + try: print_all_stacks() - self.on_detected() + except Exception: # noqa: BLE001 - stack dump must not block hard kill + pass + + # Set _detected last so observers (and tests) see it only once + # diagnostics are done and on_detected is about to fire. + with self.lock: + self._detected = True + self.on_detected() def detected(self): """Return True if hang is detected.""" diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 7023745e7dd4..ec9dbae45d7c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -1661,6 +1661,54 @@ def _copy_scratch_metadata_to_device( self._device_scratch_slots_staging, ) + def _stage_host_block_offsets_for_copy(self) -> torch.Tensor: + """Snapshot the persistent host block-offset buffer into a fresh pinned + buffer to serve as the private source of an asynchronous H2D copy. + + Same bug and same fix as nvbug 6293536 in KVCacheManager + (``_stage_block_offsets_for_copy``): the persistent buffer is rewritten + whenever blocks are allocated or freed, which happens at *scheduling* + time -- and with the overlap scheduler that is while the previous + step's H2D from this buffer can still be queued behind an in-flight + forward. The queued copy then reads the next step's block layout and + the attention kernels walk freed or reassigned blocks. + + Lifetime: for the per-layer branch the consumer is a torch ``copy_``, + which records the pinned block's use with the caching host allocator, + so dropping the Python reference is safe. The flat branch hands the + snapshot to a C++ op torch knows nothing about -- its caller must keep + the snapshot referenced until the op's stream has drained past it + (see ``_retain_flat_snapshot``). + """ + src = self.host_kv_cache_block_offsets + host_block_offsets = torch.empty( + src.shape, + dtype=src.dtype, + device="cpu", + pin_memory=prefer_pinned(), + ) + host_block_offsets.copy_(src) + return host_block_offsets + + def _retain_flat_snapshot(self, snapshot: torch.Tensor) -> None: + """Keep a pinned snapshot alive until the C++ copy op has consumed it. + + The op reads the host buffer at *execution* time on ``self._stream``, + and torch records no event for it, so the caching host allocator would + otherwise hand the block to the next same-size pinned allocation while + the op is still queued. Hold (event, snapshot) pairs and drop the ones + whose event has completed -- no synchronization, bounded by the + scheduler's run-ahead depth. + """ + retained = getattr(self, "_flat_snapshots_in_flight", None) + if retained is None: + retained = self._flat_snapshots_in_flight = [] + event = torch.cuda.Event() + with torch.cuda.stream(self._stream): + event.record() + retained.append((event, snapshot)) + retained[:] = [(e, t) for e, t in retained if not e.query()] + def _copy_batch_block_offsets_per_layer( self, dst_tensor: torch.Tensor, @@ -1670,8 +1718,9 @@ def _copy_batch_block_offsets_per_layer( num_seqs: int, ) -> None: device_copy_idx = self._copy_idx_to_device(copy_idx) + host_block_offsets = self._stage_host_block_offsets_for_copy() self._device_kv_cache_block_offsets_input.copy_( - self.host_kv_cache_block_offsets, + host_block_offsets, non_blocking=True, ) scratch_begs, scratch_ends, scratch_slots = self._copy_scratch_metadata_to_device( @@ -3535,14 +3584,33 @@ def update_resources( if not kv_cache.is_active: continue rewind_len = req.py_rewind_len + # Tokens this step actually claims: rejections inside the verify + # window plus the accepted ones. + runtime_draft_len = req.py_rewind_len + req.py_num_accepted_draft_tokens if self.is_draft: - runtime_draft_len = req.py_rewind_len + req.py_num_accepted_draft_tokens # Dynamic-tree draft managers reserve K * max_draft_len slots, # which can exceed the tree's runtime draft width. Reclaim that # reserve slack together with rejected draft tokens; otherwise # it accumulates in the draft KV cache after every generation - # step. Target managers do not allocate this reserve slack. - rewind_len += max(self._kv_reserve_draft_tokens - runtime_draft_len, 0) + # step. + reserved = self._kv_reserve_draft_tokens + else: + # The target manager reserves for the FULL drafted block -- + # prepare_resources adds `max(draft_len, reserve)` tokens every + # generation step -- but ragged verification only ever writes + # `verified_len` of them, and py_rewind_len covers rejections + # *within that window* only. The positions the scheduler trimmed + # away were reserved and never reclaimed, leaking + # (draft_len - verified_len) tokens per step until the sequence + # demands more blocks than max_seq_len allows and the KV cache + # raises "User-provided base page indices is too short". + # + # Uniform speculation is unaffected: there verified_len == + # draft_len, so this term is zero and the arithmetic is exactly + # what it was. + reserved = max(get_draft_token_length(req), + self._kv_reserve_draft_tokens) + rewind_len += max(reserved - runtime_draft_len, 0) new_capacity = ( None if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) @@ -3581,14 +3649,16 @@ def copy_batch_block_offsets( ) return + host_block_offsets = self._stage_host_block_offsets_for_copy() copy_batch_block_offsets_to_device( - self.host_kv_cache_block_offsets, + host_block_offsets, dst_tensor, copy_idx, self.index_scales, self.kv_offset, self._stream.cuda_stream, ) + self._retain_flat_snapshot(host_block_offsets) @staticmethod def _derive_reuse_salt(cache_salt: str | None) -> int | None: diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 2b1b63725fb8..469384fa3713 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1265,3 +1265,22 @@ def get_draft_token_length(request: LlmRequest) -> int: if request.py_draft_tokens is not None: return len(request.py_draft_tokens) return 0 + + +def get_request_tokens_per_gen_step(request: LlmRequest, + runtime_tokens_per_gen_step: int) -> int: + """How many tokens this request contributes to one generation step. + + Uniform speculation gives every request the batch-wide + ``runtime_tokens_per_gen_step`` (1 accepted token + ``runtime_draft_len`` + draft positions). Ragged verification -- where a scheduler hands each + request its own verify window -- records the per-request draft length in + ``py_verify_len``. + + ``py_verify_len`` is set only by the ragged path, so every other + speculation mode gets the batch-wide value back unchanged. + """ + verify_len = getattr(request, "py_verify_len", None) + if verify_len is None: + return runtime_tokens_per_gen_step + return 1 + int(verify_len) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 694be1a14a74..fd3ca223d424 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -67,6 +67,7 @@ restore_attn_metadata_after_draft_replay, update_spec_config_from_loaded_model) from ..speculative.drafting_loops import BaseDraftingLoopWrapper +from ..speculative.dspark_ragged import ragged_gather_index_lists from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..utils import (get_model_extra_attrs, @@ -81,7 +82,8 @@ from .kv_cache_manager_v2 import KVCacheManagerV2 from .layerwise_nvtx_marker import LayerwiseNvtxMarker from .llm_request import (LlmRequest, LlmRequestState, get_draft_token_length, - get_multimodal_embedding_lengths) + get_multimodal_embedding_lengths, + get_request_tokens_per_gen_step) from .mamba_cache_manager import MambaHybridCacheManager from .model_loader import ModelLoader, _construct_checkpoint_loader from .resource_manager import (BaseResourceManager, KVCacheManager, @@ -599,6 +601,9 @@ def __init__( self.is_warmup = False self.previous_request_ids = [] + # Per-request verify windows seen by the last full _prepare_tp_inputs + # pass; all None without ragged verification. + self.previous_verify_lens = [] self.has_previous_device_draft = False self.previous_accepted_tokens_cuda = torch.empty((self.batch_size, ), dtype=torch.int, @@ -697,6 +702,23 @@ def __init__( self.previous_batch_indices_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') + # Ragged verification windows, read *inside* the captured graph: they + # must live at a stable address and be written in place. + self.ragged_verify_lens_cuda = torch.empty((self.batch_size + 1, ), + dtype=torch.int, + device='cuda') + self.ragged_qo_indptr_cuda = torch.empty((self.batch_size + 2, ), + dtype=torch.int, + device='cuda') + # cap-accept: per-request acceptance ceilings for an untrimmed batch + # and the positions they discarded. Same stable-address requirement as + # above; nothing here ever syncs. + self.accept_caps_cuda = torch.empty((self.batch_size + 1, ), + dtype=torch.int, + device='cuda') + self.accept_cap_trim_cuda = torch.zeros((self.batch_size + 1, ), + dtype=torch.int, + device='cuda') self.input_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') @@ -792,6 +814,12 @@ def __init__( ) self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) + # Pinned staging buffers for async H2D copies on the ragged path; + # see _pinned_host for the two-slot WAR-guard rationale. + self._pinned_host_cache = {} + self._pinned_host_events = {} + self._pinned_host_active = {} + # Create Encoder CUDA graph config and runner. encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( use_cuda_graph=(self._is_encode_only @@ -1794,6 +1822,46 @@ def _get_graphs_to_capture( f"Capturing {len(graphs)} graphs: {graphs}") return graphs + # Case 2b: DSpark confidence-scheduled verification. The drafted block + # is always full, so draft_len is pinned to the top tier in every + # mode; `compact` moves the tiers to the token axis, `cap-accept` and + # the static fallback keep one graph per batch size. Captured graphs + # cost KV cache, so capture no tier a mode cannot select. + if (self.spec_config is not None and getattr( + self.spec_config, "enable_confidence_scheduling", False)): + tiers = sorted({ + int(t) + for t in self.spec_config.verify_len_tiers if int(t) >= 0 + }) + # The MODE, not the config flag: `cap-accept` also sets + # enable_ragged_verify=True but submits the FULL block, so it must + # not capture the ragged token ladder. + from ..speculative.dspark_observability import \ + trims_submitted_tokens + if trims_submitted_tokens(self.spec_config): + # One token bucket per tier -- exactly the totals the uniform + # ladder would produce, so the graph count does not grow. + top_tier = tiers[-1] + graphs = [(bs, top_tier, bs * (t + 1)) + for bs in cuda_graph_batch_sizes for t in tiers] + logger.info( + f"DSpark ragged verification: capturing {len(graphs)} graphs " + f"({len(cuda_graph_batch_sizes)} batch sizes x {len(tiers)} " + f"token buckets from tiers {tiers}; draft_len pinned to " + f"{top_tier}).") + return graphs + # One draft length, not the ladder: nothing varies draft_len at + # runtime, so lower tiers would be captured graphs never replayed. + top_tier = tiers[-1] + graphs = [(bs, top_tier) for bs in cuda_graph_batch_sizes] + logger.info( + f"DSpark confidence scheduling: capturing {len(graphs)} graphs " + f"({len(cuda_graph_batch_sizes)} batch sizes, draft_len pinned " + f"to {top_tier}). The drafted block is always full; only " + f"verification is trimmed, and cap-accept trims it at " + f"acceptance rather than in the layout.") + return graphs + # Case 3: Target model (two-model) or one-model without dynamic draft # Match the runtime_draft_len semantics enforced in _prepare_tp_inputs: # logical K for linear-tree modes, total tree tokens for tree decoding. @@ -1938,7 +2006,9 @@ def prepare_cross_batch(batch: ScheduledRequests, request._cached_tokens = cached_tokens request._cached_tokens_set = cached_tokens_set - def _run_capture_pass(force_non_greedy: bool, label: str) -> None: + def _run_capture_pass(force_non_greedy: bool, + label: str, + entries=None) -> None: spec_metadata = self.spec_metadata if force_non_greedy and spec_metadata is not None: spec_metadata._force_non_greedy_for_capture = True @@ -1948,16 +2018,31 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: # in this pass uses the non-greedy key; populate's override # below will keep it False on every subsequent iteration. spec_metadata.is_all_greedy_sample = False + elif spec_metadata is not None: + # Symmetric pre-flip for interleaved capture: a greedy capture + # right after an advanced one would otherwise key on the stale + # False before populate corrects it. + spec_metadata.is_all_greedy_sample = True + _graph_count = 0 try: - for bs, draft_len in graphs_to_capture: + for entry in (graphs_to_capture + if entries is None else entries): + # Ragged verification adds a token-count axis, so entries + # may be (bs, draft_len) or (bs, draft_len, token_bucket). + bs, draft_len = entry[0], entry[1] + verify_bucket = entry[2] if len(entry) > 2 else None if bs > self.batch_size: continue for max_seq_len in max_seq_len_list: + _graph_count += 1 warmup_request = self._create_cuda_graph_warmup_request( resource_manager, bs, draft_len, max_seq_len) with self._release_batch_context( warmup_request, resource_manager) as batch: + if batch is not None and verify_bucket is not None: + self._set_warmup_ragged_windows( + batch, verify_bucket, draft_len) if batch is None: # No KV cache space for this batch size. During KV # cache estimation this makes the profiling peak @@ -1987,7 +2072,31 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: finally: if force_non_greedy and spec_metadata is not None: spec_metadata._force_non_greedy_for_capture = False - + # Warmup published a bucket per captured shape; the first real + # step must not inherit the last one (a step that never reaches + # the fit would otherwise key on a capture-time bucket). + self.cuda_graph_runner.agreed_ragged_bucket = None + # Absolute private-pool size, not a per-pass delta (passes + # share the allocator); CUDA graphs live in private pools + # (segment_pool_id != (0, 0)). + _pool_bytes = 0 + _pool_ids = set() + for _seg in torch.cuda.memory_snapshot(): + _pid = _seg.get("segment_pool_id", (0, 0)) + if _pid and tuple(_pid) != (0, 0): + _pool_bytes += _seg.get("total_size", 0) + _pool_ids.add(tuple(_pid)) + logger.info( + f"CUDA graph capture: {_graph_count} graphs this pass, " + f"private pools now {_pool_bytes / (1 << 30):.3f} GiB " + f"across {len(_pool_ids)} pool(s) " + f"(ragged={getattr(self.spec_config, 'enable_ragged_verify', False)})") + + # The capture loop stamps `self.runtime_draft_len` per captured shape + # and nothing puts it back, so the engine would exit capture claiming + # the LAST shape's draft length; restore it after the passes, before + # the generic token warmup builds requests at the full draft length. + saved_runtime_draft_len = self.runtime_draft_len # Pass 1: greedy fast-path (dummy requests carry no sampling params, # so is_all_greedy_sample is naturally True). _run_capture_pass(force_non_greedy=False, label="greedy") @@ -2004,6 +2113,7 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: _run_capture_pass(force_non_greedy=True, label="advanced sampling") # Set the value back to the original value after cuda graph warmups are complete self.enable_spec_decode = self.is_spec_decode + self.runtime_draft_len = saved_runtime_draft_len # The advanced-sampling capture pass above leaves is_all_greedy_sample # set to False on spec_metadata. Reset it to the default so the first # real iteration's graph-key selection is not seeded with this @@ -2974,14 +3084,10 @@ def _set_spec_metadata_all_rank_num_tokens( # sequence contributes one token per iteration). spec_metadata.all_rank_num_tokens = spec_all_rank_num_tokens spec_metadata.all_rank_num_seqs = all_rank_num_seqs - # DSpark can draft only after the target processes the current bonus token, - # because it consumes captured target-layer hidden states for that token. - # Prefill computes hidden states for prompt tokens; the first generated token - # is sampled from the last prompt logits and has not itself passed through the - # target layers. Thus context requests seed the rolling window but do not run - # the draft. On mixed steps, num_seqs therefore over-counts the draft MoE - # workload; gen-only per-rank counts keep the FUSED_COMM (DeepGEMM MegaMoE) - # chunk loop identical across EP ranks. + # DSpark drafts only for generation requests (it needs the bonus + # token's target hidden states), so on mixed steps num_seqs + # over-counts the draft MoE workload; gen-only per-rank counts keep + # the FUSED_COMM chunk loop identical across EP ranks. if all_rank_num_gens is not None: spec_metadata.all_rank_num_gens = all_rank_num_gens if (spec_metadata.spec_dec_mode.is_mtp_eagle_one_model() @@ -3265,6 +3371,16 @@ def _can_use_incremental_update( if self.previous_request_ids != request_ids: return False + # The incremental path is only valid while every request's sequence + # length is unchanged; ragged verification re-picks windows every + # step, so fall back to a full prepare whenever they moved. + verify_lens = [ + getattr(request, "py_verify_len", None) + for request in scheduled_requests.generation_requests + ] + if self.previous_verify_lens != verify_lens: + return False + has_current_device_draft = next_draft_tokens_device is not None return (self.is_draft_model and self.model_is_wrapped) or ( has_current_device_draft and self.has_previous_device_draft) @@ -3503,18 +3619,604 @@ def _apply_incremental_update_draft( return inputs, self.gather_ids_cuda[:num_generation_tokens] + def _set_warmup_ragged_windows(self, batch, verify_bucket: int, + draft_len: int) -> None: + """Give a capture batch per-request windows summing to ``verify_bucket``. + + The warmup batch must hit the bucket total exactly, or the graph is + keyed under a total no runtime batch produces and every ragged step + silently falls back to eager. Any split with the right sum captures + the right shape (the raggedness lives in the contents of qo_indptr). + """ + requests = batch.generation_requests + if not requests: + return + from ..speculative.dspark_ragged import RaggedVerifyLayout + + lens = torch.ones(len(requests), dtype=torch.int32) + layout = RaggedVerifyLayout.from_verify_lens( + lens, + graph_num_tokens=int(verify_bucket), + total_verify_tokens=len(requests)) + filled = layout.fill_bucket(max_verify_len=1 + int(draft_len)) + for request, tokens in zip(requests, filled.verify_lens.tolist()): + request.py_verify_len = int(tokens) - 1 + # Publish to the runner too: the graph key carries this value, and the + # runtime fit never runs during capture. + self.cuda_graph_runner.agreed_ragged_bucket = int(verify_bucket) + + def ragged_verify_token_buckets(self, padded_bs: int) -> List[int]: + """Captured token totals for a ragged batch of ``padded_bs`` rows. + + One bucket per verify-length tier -- exactly the totals the uniform + ladder would produce, so the captured graph set does not grow. + """ + if self.spec_config is None: + return [] + tiers = sorted({ + int(t) + for t in self.spec_config.verify_len_tiers + }) + return [int(padded_bs) * (t + 1) for t in tiers] + + def _record_dspark_graph_use(self, *, replayed: bool, shape=None) -> None: + """Count graph replay vs eager for the DSpark ragged stats, if enabled.""" + stats = self._dspark_ragged_stats + if stats is not None: + stats.record_graph(replayed=replayed, shape=shape) + + @property + def _dspark_ragged_stats(self): + """The DSpark worker's stats object, or None on every other path.""" + worker = self._get_spec_worker() + return getattr(worker, "ragged_stats", None) if worker else None + + def fit_ragged_verify_lens( + self, + generation_requests, + verify_lens: List[int], + peer_stats: Optional[List[List[int]]] = None) -> Optional[int]: + """Round this batch's verify total onto a captured bucket, in place. + + The flat token count is a CUDA-graph key axis, so the total must land + on a captured value *before* the key is built; the rounding slack is + spent on real verification (``RaggedVerifyLayout.fill_bucket``). + Writes the fitted window onto each request's ``py_verify_len`` and + returns the bucket, or None when no captured shape fits (the caller + falls back to uniform scheduling). + + ``peer_stats`` is every attention-DP rank's ``[num_requests, + total_tokens, all_can_graph]``, including this rank's: the bucket must + be sized from the cross-rank maximum or ranks pick different shapes, + and ``all_can_graph`` is the group's answer to whether that maximum + will be taken at all. + """ + runner = self.cuda_graph_runner + # Cleared on entry so every early return leaves no stale value: a + # bucket carried over from a previous step would key into the wrong + # graph. The padded row count follows the same rule. + runner.agreed_ragged_bucket = None + self._dspark_last_padded_bs = None + if not runner.enabled or not verify_lens: + return None + if len(verify_lens) != len(generation_requests): + # A partially windowed batch would go ragged in the token layout + # while the spec metadata stayed uniform; refuse instead. + logger.debug( + f"DSpark ragged: got {len(verify_lens)} verify lengths for " + f"{len(generation_requests)} generation requests; falling back " + f"to uniform scheduling") + return None + tiers = sorted({ + int(t) + for t in self.spec_config.verify_len_tiers + }) + if not tiers: + return None + max_verify_len = 1 + tiers[-1] + token_lens = [1 + int(v) for v in verify_lens] + + from ..speculative.dspark_ragged import (RaggedVerifyLayout, + choose_ragged_capture_shape) + + bs_buckets = sorted({int(b) for b in runner.supported_batch_sizes}) + if not bs_buckets: + return None + # Resolve rows first, then tokens: the bucket grid depends on the + # widest rank's row count. `all_can_graph` (third peer-stat element; + # absent means single-rank) says whether _get_padded_batch will take + # the cross-rank row maximum at all -- only when EVERY rank can graph. + all_can_graph = (bool(peer_stats[0][2]) if peer_stats + and len(peer_stats[0]) > 2 else True) + if not all_can_graph: + # No rank will replay a graph this step, and _get_padded_batch + # pads differently than this fit assumes on such steps, so the + # budget would go to rows that never appear. Decline (conservative: + # an eager ragged step would still save compute). + logger.debug( + "DSpark ragged: some rank cannot run a CUDA graph this step, " + "so no captured bucket applies; falling back to uniform") + return None + widest_rows = max([int(s[0]) for s in peer_stats], + default=len(token_lens)) if peer_stats else len( + token_lens) + padded_bs = runner._round_up_batch_size(widest_rows) + if padded_bs == 0: + return None + # The fit assumes the batch will actually be padded to `padded_bs`; + # padding may decline, and then the fitted total lands in no captured + # bucket. Decline to go ragged rather than fit against rows that will + # not exist (`will_pad_to` mirrors only the SIZE guards; the group's + # graph answer is `all_can_graph` above). + if not runner.will_pad_to(padded_bs, len(token_lens)): + logger.debug( + f"DSpark ragged: padding to {padded_bs} rows is not available " + f"for {len(token_lens)} requests, so the bucket grid derived " + f"from it would not be realised; falling back to uniform") + return None + buckets = self.ragged_verify_token_buckets(padded_bs) + if not buckets: + return None + try: + shape = choose_ragged_capture_shape( + num_real_requests=len(token_lens), + total_verify_tokens=sum(token_lens), + bs_buckets=bs_buckets, + token_buckets=buckets, + peer_stats=peer_stats, + # With the window cap the chooser also enforces per-rank + # decomposability from the allgathered stats, so accept vs + # decline is identical on every rank -- the local checks + # below become unreachable safety nets instead of the + # group-splitting decision points they used to be. + max_verify_len=max_verify_len) + except ValueError as exc: + logger.debug(f"DSpark ragged shape selection failed, falling back " + f"to uniform scheduling: {exc}") + return None + padded_bs, bucket = shape.padded_bs, shape.bucket + + # Pad rows are appended later as a *single shared dummy object*, so + # they all carry the same window and their contribution must be + # decided here, not left to the fill. They take the minimum (one + # token) so real requests get the slack, growing only when the real + # rows cannot absorb the bucket. + n_real = len(token_lens) + n_pad = padded_bs - n_real + real_capacity = n_real * max_verify_len + floor_tokens = sum(token_lens) + if n_pad == 0: + pad_len = 1 + else: + # pad_len has to leave the real rows a target they can actually + # hit: at least what they already need, at most their capacity. + # bucket - n_pad*pad_len >= floor_tokens -> upper bound + # bucket - n_pad*pad_len <= real_capacity -> lower bound + lo = max(1, -(-(bucket - real_capacity) // n_pad)) # ceil division + hi = min(max_verify_len, (bucket - floor_tokens) // n_pad) + if lo > hi: + # Unreachable when the chooser ran with max_verify_len (it + # verified this rank's decomposition from the payload); a hit + # here means the two arithmetics diverged -- worth a warning, + # not a silent debug line. + logger.warning( + f"DSpark ragged: bucket {bucket} admits no pad-row window " + f"for {n_real} real ({floor_tokens}..{real_capacity} " + f"tokens) + {n_pad} pad rows; falling back to uniform " + f"scheduling") + return None + pad_len = lo + real_target = bucket - n_pad * pad_len + if real_target < floor_tokens or real_target > real_capacity: + logger.warning( + f"DSpark ragged: bucket {bucket} leaves {real_target} tokens " + f"for {n_real} real requests ({n_pad} pad rows at {pad_len}), " + f"outside [{floor_tokens}, {real_capacity}]; falling back " + f"to uniform scheduling") + return None + + lens_tensor = torch.tensor(token_lens, dtype=torch.int32) + layout = RaggedVerifyLayout.from_verify_lens( + lens_tensor, + graph_num_tokens=real_target, + total_verify_tokens=sum(token_lens)) + try: + filled = layout.fill_bucket(max_verify_len=max_verify_len) + except ValueError as exc: + # A mismatch means this batch has no captured shape; degrade to + # uniform rather than run eager. + logger.debug(f"DSpark ragged bucket fit failed, falling back to " + f"uniform scheduling: {exc}") + return None + + # Published only now, past the last way this fit can fail: a stale + # published bucket on a fallback step is the state behind the ragged + # IMA. `_dspark_last_padded_bs` feeds the + # stats; `ragged_pad_verify_len` is stamped on the shared dummy by + # _get_padded_batch. + self._dspark_last_padded_bs = int(padded_bs) + self._dspark_last_num_real = int(n_real) + runner.agreed_ragged_bucket = int(bucket) + runner.ragged_pad_verify_len = pad_len - 1 + + published = filled.verify_lens.tolist() + for request, tokens in zip(generation_requests, published): + request.py_verify_len = int(tokens) - 1 + return int(bucket) + + def _apply_device_window_prologue(self, inputs, + new_tensors_device) -> bool: + """Re-rank this step's verify windows on device, with fresh confidence. + + Runs after ``_prepare_inputs`` and before the graph replay, all device + ops on the current stream -- stream order makes every write visible to + the replayed graph, and nothing here reads device data back to the + host. The host has already agreed ``(padded_bs, bucket)`` and staged a + SHAPE SPLIT of the (lagged) budget through the normal fit; this + prologue re-distributes the same real/pad token totals by the verified + block's OWN confidence (ranking lag zero) and overwrites the layout's + content: verify lens, qo_indptr, the per-token row maps, the packed + input/position/draft tokens, and the per-request kv-length delta from + the shape split to the true windows. + + Returns True when applied; False when a precondition fails (the step + then runs the shape split as-is, which is a valid window assignment). + """ + from ..speculative.dspark_device_select import select_windows_device + + runner = self.cuda_graph_runner + budget = self._dspark_device_budget + self._dspark_device_budget = None + if budget is None or runner.agreed_ragged_bucket is None: + return False + if not getattr(self, "_dspark_prev_covers_batch", False): + # The input gathers below address new_tokens_device by + # previous_batch_indices in batch order; a request without a + # previous device tensor breaks that addressing. + return False + worker = self._get_spec_worker() + planner = worker.verify_planner + if planner is None or worker.staged_confidence_buffer() is None: + return False + if worker.batch_slot_view(1) is None: + return False + spec_metadata = inputs.get('spec_metadata') + attn_metadata = inputs.get('attn_metadata') + if spec_metadata is None or attn_metadata is None: + return False + if int(attn_metadata.num_contexts) != 0: + return False + + bucket = int(runner.agreed_ragged_bucket) + padded_bs = int(self._dspark_last_padded_bs) + n_real = int(self._dspark_last_num_real) + pad_len_tok = int(runner.ragged_pad_verify_len) + 1 + real_tokens = bucket - (padded_bs - n_real) * pad_len_tok + cfg = planner.cfg + # The scheduler cannot grant more than the real rows can absorb under + # the published split; the fill tops up any shortfall. + budget = max( + 0, + min(int(budget), real_tokens - n_real * (cfg.min_verify_len + 1))) + + # Snapshot the shape split BEFORE overwriting: past_seen per row is + # the staged position at each row's first token, and the kv delta + # needs the split the host baked into kv_lens_cuda. + lens_buf = self.ragged_verify_lens_cuda + qo_buf = self.ragged_qo_indptr_cuda + split_lens = lens_buf[:padded_bs].clone() + split_qo = qo_buf[:padded_bs + 1].to(torch.long) + past_seen = self.position_ids_cuda[split_qo[:-1]].clone() + + expected_stamp = worker.verified_draft_seq_cuda() + stamps = worker.confidence_stamp_buffer( + ) if expected_stamp is not None else None + result = select_windows_device( + confidence_logits=worker.staged_confidence_buffer(), + slot_idx=worker.batch_slot_view(padded_bs), + num_real=n_real, + budget=budget, + graph_num_tokens=bucket, + cfg=cfg, + apply_calibration=planner.apply_calibration, + stamp=stamps, + expected_stamp=expected_stamp, + pad_len=pad_len_tok, + ) + + lens_buf[:padded_bs].copy_(result.verify_lens) + qo_buf[:padded_bs + 1].copy_(result.qo_indptr) + # The host stages kv_lens as num_cached + seq_lens_kv, and BOTH terms + # bake the per-request token window (num_cached = past + tokens; + # seq_lens_kv = tokens), so kv_lens = past + 2*S -- the window counts + # TWICE. Moving to the true windows therefore needs 2*(w - S), not + # (w - S): the single-delta variant left every re-ranked request's + # kv_len off by (w - S), which shifted the indexer K-cache slot + # mapping (slot_mapping_fp8) and silently wrote K entries into the + # wrong cache slots. Established empirically by an A/B tensor diff + # against a full host restage with the same windows. + window_delta = result.verify_lens - split_lens + attn_metadata.kv_lens_cuda[:padded_bs] += 2 * window_delta + # The graph adds previous_kv_lens_offsets (staged as new_tokens_lens - + # shape_lens, per request) to kv_lens during replay; host-with-w + # stages new_tokens_lens - w there, so the offsets move by -(w - S). + # Combined: (past + 2S) + 2(w-S) + (new - S) - (w-S) = past + w + new, + # exactly the host-with-w in-graph sum. + self.previous_kv_lens_offsets_cuda[:padded_bs] -= window_delta.to( + self.previous_kv_lens_offsets_cuda.dtype) + + req_idx = result.req_idx + device = req_idx.device + flat = torch.arange(bucket, device=device) + offset = flat - result.qo_indptr.to(torch.long)[req_idx] + prev_slots = self.previous_batch_indices_cuda[:n_real].to(torch.long) + slots_tok = prev_slots[req_idx.clamp(max=n_real - 1)] + + new_tokens_device = new_tensors_device.new_tokens + new_tokens_lens_device = new_tensors_device.new_tokens_lens + next_draft_tokens_device = new_tensors_device.next_draft_tokens + + self.input_ids_cuda[:bucket] = new_tokens_device.transpose( + 0, 1)[slots_tok, offset].flatten().to(self.input_ids_cuda.dtype) + self.position_ids_cuda[:bucket] = past_seen[req_idx] + offset.to( + past_seen.dtype) + # Overlap corrections gather by each token's OWNER; rebuild both the + # index and the per-token offset it feeds (mirrors the host staging at + # the previous_pos_indices block). + self.previous_pos_indices_cuda[:real_tokens] = slots_tok[:real_tokens].to( + self.previous_pos_indices_cuda.dtype) + self.previous_pos_id_offsets_cuda[:real_tokens].copy_( + new_tokens_lens_device[slots_tok[:real_tokens]]) + # Draft tokens pack compactly: token t at in-window offset o >= 1 goes + # to slot t - owner - 1 (one bonus token per preceding row). Anchor + # tokens (o == 0) are parked one past the real draft region, where any + # valid token id is discard-territory. + real_draft = real_tokens - n_real + mask = offset >= 1 + dst = torch.where(mask, flat - req_idx - 1, + torch.full_like(req_idx, real_draft)) + src = next_draft_tokens_device[slots_tok, (offset - 1).clamp_min(0)] + self.draft_tokens_cuda[dst] = src.to(self.draft_tokens_cuda.dtype) + + attn_metadata.apply_device_ragged_layout(result.verify_lens, req_idx, + result.kv_correction) + return True + + @staticmethod + def _ragged_token_lens(generation_requests) -> Optional[List[int]]: + """Each generation request's token window, or None if the batch is uniform. + + ``py_verify_len`` counts drafted positions; the token window adds the + bonus position. None unless *every* request carries a window -- + partially windowed batches are silent token misattribution. + """ + verify_lens = [ + getattr(request, "py_verify_len", None) + for request in generation_requests + ] + if not verify_lens or any(v is None for v in verify_lens): + return None + return [1 + int(v) for v in verify_lens] + + def _publish_gen_token_layout(self, attn_metadata, + generation_requests) -> None: + """Hand the attention metadata this step's gen-token layout, before it + prepares. + + Split out of :meth:`_attach_ragged_verify_layout` for ordering: + ``attn_metadata.prepare()`` is the only consumer of + ``ragged_verify_lens`` and runs well before the spec metadata is + assembled. The uniform stride is published here too, because + ``max_draft_tokens`` is a static buffer-sizing ceiling that does not + move when a shorter tier is chosen. It cannot live in + ``update_spec_dec_param`` either: that runs against the base metadata, + and ``prepare()`` runs against the per-key CUDA-graph copy. + """ + if attn_metadata is None: + return + if hasattr(attn_metadata, "runtime_tokens_per_gen_step"): + attn_metadata.runtime_tokens_per_gen_step = ( + self.get_runtime_tokens_per_gen_step(self.runtime_draft_len)) + if hasattr(attn_metadata, "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = self._ragged_token_lens( + generation_requests) + if hasattr(attn_metadata, "device_windows_mode"): + # Tells the attention prepare that the host window VALUES are a + # shape split (bounds only); the true windows land on device + # through apply_device_ragged_layout after prepare. + from ..speculative.dspark_verify import device_windows_enabled + attn_metadata.device_windows_mode = device_windows_enabled() + + def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: + """A persistent pinned staging buffer holding ``values``. + + Async H2D sources must outlive the queued copy -- PyTorch does not + extend the lifetime of an async source, so a temporary tensor can be + reclaimed while the DMA still reads it. Buffers are keyed by name and + grown monotonically. + """ + values = list(values) + # WAR guard: a slot must not be rewritten while its previous + # non_blocking H2D is still queued. Callers record the active slot's + # event via _pinned_host_record after enqueuing; two slots alternate, + # so this wait targets the copy from two steps ago and is free in + # steady state (a single slot would stall prepare behind the previous + # step's stream position). + slot = 1 - self._pinned_host_active.get(key, 1) + evt = self._pinned_host_events.get((key, slot)) + if evt is not None: + evt.synchronize() + bufs = self._pinned_host_cache.setdefault(key, [None, None]) + buf = bufs[slot] + if buf is None or buf.numel() < len(values) or buf.dtype != dtype: + buf = torch.empty(max(len(values), 1), + dtype=dtype, + pin_memory=prefer_pinned()) + bufs[slot] = buf + self._pinned_host_active[key] = slot + view = buf[:len(values)] + if values: + view.copy_(torch.tensor(values, dtype=dtype)) + return view + + def _pinned_host_record(self, *keys: str) -> None: + """Mark the H2D copies from these staging buffers as enqueued. + + Call after the non_blocking copy that reads a `_pinned_host` view. + Skipped during graph capture (the staging copies are eager on every + step, including graph steps; only the one-time capture pass is not). + """ + if torch.cuda.is_current_stream_capturing(): + return + for key in keys: + slot = self._pinned_host_active.get(key) + if slot is None: + continue + evt = self._pinned_host_events.get((key, slot)) + if evt is None: + evt = self._pinned_host_events[(key, slot)] = torch.cuda.Event() + evt.record() + + def _attach_accept_caps(self, spec_metadata, generation_requests) -> None: + """Publish ``cap-accept`` ceilings, without touching the layout. + + ``py_verify_cap`` is kept separate from ``py_verify_len`` so every + layout consumer (input packing, DSA row expansion, graph key, KV + rewind) keeps seeing an ordinary uniform batch. The cap is in + ``num_accepted_tokens`` units (accepted drafts plus the bonus), hence + ``1 + py_verify_cap``. + """ + caps = [ + getattr(request, "py_verify_cap", None) + for request in generation_requests + ] + worker = self._get_spec_worker() + mode = getattr(worker, "ragged_verify_mode", None) if worker else None + cap_accept_mode = (mode is not None and mode.computes_windows + and not mode.trims_submitted_tokens) + if not caps or any(cap is None for cap in caps): + if cap_accept_mode and caps: + # cap-accept mode: NEVER publish None. `apply_accept_caps` + # reads `spec_metadata.accept_caps` as a Python attribute at + # CAPTURE time, so a None on the capture step (warmup batches + # carry no caps) bakes a no-op into the graph and every replay + # silently skips the clamp -- the zero-trim positive-control + # failure. Capless requests get the full block: numerically a + # no-op clamp, but the ops and buffer address stay captured. + caps = [ + self.spec_config.max_draft_len if cap is None else cap + for cap in caps + ] + else: + # Any other mode: partial caps are meaningless; drop them all, + # but warn, because the step then silently runs as `static`. + if caps and any(cap is not None for cap in caps): + if not getattr(self, "_warned_partial_accept_caps", False): + self._warned_partial_accept_caps = True + logger.warning( + "DSpark cap-accept: %d of %d generation requests " + "carry no verify cap, so caps were dropped for " + "the whole batch and this step ran as 'static'. " + "The mode is measuring nothing on steps like " + "this one.", + sum(1 for cap in caps if cap is None), len(caps)) + spec_metadata.accept_caps = None + spec_metadata.accept_cap_trim = None + return + + n = len(caps) + # Persistent buffer: acceptance runs inside the captured graph, so a + # fresh tensor here would be invisible to every replay. + caps_view = self.accept_caps_cuda[:n] + caps_view.copy_(self._pinned_host("accept_caps", + [1 + int(cap) for cap in caps], + torch.int32), + non_blocking=True) + self._pinned_host_record("accept_caps") + spec_metadata.accept_caps = caps_view + spec_metadata.accept_cap_trim = self.accept_cap_trim_cuda + + def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, + generation_requests) -> None: + """Publish this step's per-request verify windows to both metadatas. + + Only the DSpark confidence scheduler sets ``py_verify_len``; every + other path leaves the fields None. Acceptance needs the windows to + slice correctly; the DSA indexer needs them to expand kv_lens/block + tables per query token. + """ + self._attach_accept_caps(spec_metadata, generation_requests) + token_lens = self._ragged_token_lens(generation_requests) + if token_lens is None: + spec_metadata.verify_lens = None + spec_metadata.qo_indptr = None + spec_metadata.total_verify_tokens = None + if attn_metadata is not None and hasattr(attn_metadata, + "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = None + return + + from ..speculative.dspark_ragged import build_qo_indptr + + n = len(token_lens) + # Persistent buffers: a captured graph baked in the address it saw at + # capture time, so a fresh tensor would be invisible to every replay. + lens_view = self.ragged_verify_lens_cuda[:n] + lens_view.copy_(self._pinned_host("ragged_verify_lens", token_lens, + torch.int32), + non_blocking=True) + self._pinned_host_record("ragged_verify_lens") + indptr_view = self.ragged_qo_indptr_cuda[:n + 1] + indptr_view.copy_(build_qo_indptr(lens_view), non_blocking=True) + spec_metadata.verify_lens = lens_view + spec_metadata.qo_indptr = indptr_view + spec_metadata.total_verify_tokens = sum(token_lens) + if attn_metadata is not None and hasattr(attn_metadata, + "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = token_lens + + def _ragged_gather_indices( + self, slots: List[int], + counts: List[int]) -> Tuple[torch.Tensor, torch.Tensor]: + """Row/column index pair for gathering a ragged block out of a + [num_slots, max_width] device tensor. + + ``slots[i]`` contributes ``counts[i]`` entries taken from columns + ``0..counts[i]-1``, concatenated in batch order. This replaces the + ``tensor[slots, :width]`` strided gather used when every request + verifies the same number of positions. + """ + rows, cols = ragged_gather_index_lists(slots, counts) + rows_dev = self._pinned_host("gather_rows", rows, + torch.long).to('cuda', non_blocking=True) + self._pinned_host_record("gather_rows") + cols_dev = self._pinned_host("gather_cols", cols, + torch.long).to('cuda', non_blocking=True) + self._pinned_host_record("gather_cols") + return (rows_dev, cols_dev) + def _update_target_input_tensors( - self, num_accepted_tokens_device: torch.Tensor, + self, + num_accepted_tokens_device: torch.Tensor, new_tokens_device: torch.Tensor, next_draft_tokens_device: torch.Tensor, - new_tokens_lens_device: torch.Tensor, previous_slots: torch.Tensor, - total_num_tokens: int, num_extend_reqeust_wo_dummy: int, + new_tokens_lens_device: torch.Tensor, + previous_slots: torch.Tensor, + total_num_tokens: int, + num_extend_reqeust_wo_dummy: int, num_tokens_per_extend_request: int, - previous_batch_draft_tokens: int): + previous_batch_draft_tokens: int, + tokens_per_extend_request: Optional[List[int]] = None, + previous_batch_slots: Optional[List[int]] = None): """ This function performs in-place updates on position_ids, num_accepted_draft_tokens, input_ids, draft_tokens, and offset tensors for speculative decoding extend context operations. + + ``tokens_per_extend_request`` carries the per-request token counts under + ragged verification; it is None (and every gather below stays strided + on ``num_tokens_per_extend_request``) whenever the batch is uniform. """ + is_ragged = tokens_per_extend_request is not None # Prepare position_ids idx_accepted_tokens = self.idx_accepted_tokens_cache[:total_num_tokens] @@ -3533,33 +4235,67 @@ def _update_target_input_tensors( # CRITICAL: Only extract the needed tokens based on num_tokens_per_extend_request # new_tokens_device shape: [batch, 1 + max_draft_len] # We need: [previous_batch, num_tokens_per_extend_request] - new_tokens = new_tokens_device.transpose( - 0, 1)[previous_slots, :num_tokens_per_extend_request].flatten() + if is_ragged: + token_rows, token_cols = self._ragged_gather_indices( + previous_batch_slots, tokens_per_extend_request) + new_tokens = new_tokens_device.transpose(0, + 1)[token_rows, + token_cols].flatten() + else: + new_tokens = new_tokens_device.transpose( + 0, 1)[previous_slots, :num_tokens_per_extend_request].flatten() self.input_ids_cuda[:total_num_tokens].copy_(new_tokens, non_blocking=True) # Prepare draft tokens - num_draft_tokens_per_extend_request = num_tokens_per_extend_request - 1 - self.draft_tokens_cuda[:previous_batch_draft_tokens].copy_( - next_draft_tokens_device[ - previous_slots, :num_draft_tokens_per_extend_request].flatten(), - non_blocking=True) + if is_ragged: + if previous_batch_draft_tokens > 0: + draft_rows, draft_cols = self._ragged_gather_indices( + previous_batch_slots, + [t - 1 for t in tokens_per_extend_request]) + self.draft_tokens_cuda[:previous_batch_draft_tokens].copy_( + next_draft_tokens_device[draft_rows, draft_cols].flatten(), + non_blocking=True) + else: + num_draft_tokens_per_extend_request = num_tokens_per_extend_request - 1 + self.draft_tokens_cuda[:previous_batch_draft_tokens].copy_( + next_draft_tokens_device[ + previous_slots, :num_draft_tokens_per_extend_request]. + flatten(), + non_blocking=True) - # Compute kv_len_offsets and update offset tensors - previous_pos_indices = previous_slots.repeat_interleave( - num_tokens_per_extend_request) + # Compute kv_len_offsets and update offset tensors. + # kv_len_offsets pairs with num_cached_tokens_per_seq (past_seen + + # tokens_this_step) computed by the caller: the two cancel to + # past_seen + accepted, so both must use the same per-request count. + if is_ragged: + tokens_per_request_device = torch.tensor( + tokens_per_extend_request, + dtype=torch.long, + pin_memory=prefer_pinned()).to(new_tokens_lens_device.device, + non_blocking=True) + # `output_size` is not optional: without it repeat_interleave with + # tensor `repeats` reads the cumulative sum back to the host -- a + # device->host sync on every ragged step. + previous_pos_indices = torch.repeat_interleave( + previous_slots, + tokens_per_request_device, + output_size=total_num_tokens) + previous_kv_len_offsets = (new_tokens_lens_device[previous_slots] - + tokens_per_request_device) + else: + previous_pos_indices = previous_slots.repeat_interleave( + num_tokens_per_extend_request) + kv_len_offsets_device = new_tokens_lens_device - num_tokens_per_extend_request + previous_kv_len_offsets = kv_len_offsets_device[previous_slots] self.previous_pos_indices_cuda[:total_num_tokens].copy_( previous_pos_indices, non_blocking=True) - kv_len_offsets_device = new_tokens_lens_device - num_tokens_per_extend_request - self.previous_pos_id_offsets_cuda[:num_extend_reqeust_wo_dummy * - num_tokens_per_extend_request].copy_( - new_tokens_lens_device[ - self. - previous_pos_indices_cuda[: - total_num_tokens]], - non_blocking=True) + self.previous_pos_id_offsets_cuda[:total_num_tokens].copy_( + new_tokens_lens_device[ + self.previous_pos_indices_cuda[:total_num_tokens]], + non_blocking=True) self.previous_kv_lens_offsets_cuda[:num_extend_reqeust_wo_dummy].copy_( - kv_len_offsets_device[previous_slots], non_blocking=True) + previous_kv_len_offsets, non_blocking=True) def _apply_incremental_update_target( self, @@ -3581,7 +4317,6 @@ def _apply_incremental_update_target( spec_config = self.spec_config num_tokens_per_extend_request = self.get_runtime_tokens_per_gen_step( self.runtime_draft_len) - runtime_draft_token_buffer_width = num_tokens_per_extend_request - 1 prompt_lengths = torch.empty(num_extend_requests, dtype=torch.int, @@ -3599,6 +4334,11 @@ def _apply_incremental_update_target( request_accepted_path = {} num_extend_dummy_requests = 0 num_previous_batch = 0 + # Per-request token counts, in previous-batch order. Stays all-equal to + # num_tokens_per_extend_request unless a ragged scheduler assigned each + # request its own verify window. + tokens_per_extend_request = [] + previous_batch_slots = [] use_extend_ctx = (self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( @@ -3609,11 +4349,17 @@ def _apply_incremental_update_target( request.py_num_accepted_draft_tokens_indices base_past_seen = request.max_beam_num_tokens - 1 + # Under ragged verification every length derived below has to come + # from this one per-request value; mixing it with the batch-wide + # count desynchronizes the flat token layout from the KV-length + # correction applied in _preprocess_inputs. + req_tokens_per_gen_step = get_request_tokens_per_gen_step( + request, num_tokens_per_extend_request) if use_extend_ctx: # We're treating the prompt lengths as context requests here, so # the prompt lens should not include the cached tokens. - prompt_lengths[idx] = num_tokens_per_extend_request + prompt_lengths[idx] = req_tokens_per_gen_step else: prompt_lengths[idx] = request.py_prompt_len @@ -3630,17 +4376,21 @@ def _apply_incremental_update_target( previous_batch_indices[ num_previous_batch] = request.py_batch_idx num_previous_batch += 1 + previous_batch_slots.append(request.py_batch_idx) + tokens_per_extend_request.append(req_tokens_per_gen_step) request.cached_tokens = (base_past_seen + - num_tokens_per_extend_request) + req_tokens_per_gen_step) num_cached_tokens_per_seq[idx] = ( - base_past_seen + num_tokens_per_extend_request - + base_past_seen + req_tokens_per_gen_step - request.py_num_compressed_tokens) request.py_batch_idx = request.py_seq_slot num_extend_reqeust_wo_dummy = num_extend_requests - num_extend_dummy_requests - total_num_tokens = num_extend_reqeust_wo_dummy * num_tokens_per_extend_request + is_ragged_gen = any(tokens != num_tokens_per_extend_request + for tokens in tokens_per_extend_request) + total_num_tokens = sum(tokens_per_extend_request) previous_slots = self.previous_batch_indices_cuda[:num_previous_batch] previous_slots.copy_(previous_batch_indices[:num_previous_batch], @@ -3649,8 +4399,8 @@ def _apply_incremental_update_target( prompt_lengths = prompt_lengths.tolist() num_cached_tokens_per_seq = num_cached_tokens_per_seq.tolist() - previous_batch_draft_tokens = (num_extend_reqeust_wo_dummy * - runtime_draft_token_buffer_width) + previous_batch_draft_tokens = (total_num_tokens - + num_extend_reqeust_wo_dummy) self._update_target_input_tensors( num_accepted_tokens_device=num_accepted_tokens_device, @@ -3661,10 +4411,16 @@ def _apply_incremental_update_target( total_num_tokens=total_num_tokens, num_extend_reqeust_wo_dummy=num_extend_reqeust_wo_dummy, num_tokens_per_extend_request=num_tokens_per_extend_request, - previous_batch_draft_tokens=previous_batch_draft_tokens) - - # Prepare spec_metadata - num_generation_tokens = num_extend_requests * num_tokens_per_extend_request + previous_batch_draft_tokens=previous_batch_draft_tokens, + tokens_per_extend_request=(tokens_per_extend_request + if is_ragged_gen else None), + previous_batch_slots=previous_batch_slots) + + # Prepare spec_metadata. Dummy requests are padded to the batch-wide + # window, so only the real requests contribute ragged token counts. + num_generation_tokens = ( + total_num_tokens + + num_extend_dummy_requests * num_tokens_per_extend_request) if spec_metadata is not None: total_draft_lens = self.max_total_draft_tokens * num_extend_requests spec_metadata.draft_tokens = self.draft_tokens_cuda[: @@ -4187,11 +4943,21 @@ def append_cross_attention_state(request: LlmRequest, # will contain previous batch indices of generation requests previous_batch_indices = [] previous_pos_indices = [] + # Token count each previous-batch request contributes, in batch order. + # Uniform speculation makes every entry runtime_tokens_per_gen_step; + # ragged verification does not, and the device-side gathers below + # switch to an index-list layout when they disagree. + previous_batch_tokens_per_request = [] + # Flat generation-token offset of the first previous-batch request, + # captured while walking the batch so the ragged layout does not have + # to assume a fixed stride. + previous_batch_token_start = None + extend_tokens_emitted = 0 runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( self.runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 for request in extend_requests: - if getattr(request, "py_needs_onehot_draft_probs", False): + if request.py_needs_onehot_draft_probs: if request.py_seq_slot is not None: padding_gen_slots.append(request.py_seq_slot) request.py_needs_onehot_draft_probs = False # consume once @@ -4206,13 +4972,18 @@ def append_cross_attention_state(request: LlmRequest, if next_draft_tokens_device is None or request.is_dummy or request.py_batch_idx is None: # get token ids, including input token ids and draft token ids. For these dummy requests, # no need to copy the token ids. + # Only the request's own window is submitted to the target; + # taking the full drafted length here would desynchronize the + # flat token layout from the spec metadata and the graph key. + num_draft_tokens = get_request_tokens_per_gen_step( + request, 1 + get_draft_token_length(request)) - 1 if not (request.is_attention_dp_dummy or request.is_cuda_graph_dummy): input_ids.append(request.get_last_tokens(0)) - input_ids.extend(request.py_draft_tokens) - draft_tokens.extend(request.py_draft_tokens) + input_ids.extend(request.py_draft_tokens[:num_draft_tokens]) + draft_tokens.extend( + request.py_draft_tokens[:num_draft_tokens]) # get other ids and lengths - num_draft_tokens = get_draft_token_length(request) past_seen_token_num = request.max_beam_num_tokens - 1 draft_lens.append(num_draft_tokens) if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( @@ -4236,6 +5007,7 @@ def append_cross_attention_state(request: LlmRequest, num_cached_tokens_per_seq.append( past_seen_token_num - request.py_num_compressed_tokens) request.cached_tokens = past_seen_token_num + extend_tokens_emitted += 1 + num_draft_tokens # update batch index request.py_batch_idx = request.py_seq_slot else: @@ -4243,33 +5015,46 @@ def append_cross_attention_state(request: LlmRequest, previous_batch_idx = request.py_batch_idx request.py_batch_idx = request.py_seq_slot - sequence_lengths.append(runtime_tokens_per_gen_step) + if previous_batch_token_start is None: + previous_batch_token_start = extend_tokens_emitted + + # Under ragged verification each request gets its own window; + # every length below has to come from the same per-request + # value, otherwise the flat token layout and the KV-length + # correction in _preprocess_inputs disagree. + req_tokens_per_gen_step = get_request_tokens_per_gen_step( + request, runtime_tokens_per_gen_step) + + sequence_lengths.append(req_tokens_per_gen_step) num_accepted_draft_tokens.append( request.py_num_accepted_draft_tokens) past_seen_token_num = request.max_beam_num_tokens - 1 - draft_lens.append(runtime_draft_token_buffer_width) + draft_lens.append(req_tokens_per_gen_step - 1) gather_ids.extend( list( range(len(position_ids), - len(position_ids) + runtime_tokens_per_gen_step))) + len(position_ids) + req_tokens_per_gen_step))) position_ids.extend( list( - range(past_seen_token_num, past_seen_token_num + - runtime_tokens_per_gen_step))) + range(past_seen_token_num, + past_seen_token_num + req_tokens_per_gen_step))) # previous tensor previous_batch_indices.append(previous_batch_idx) previous_pos_indices.extend([previous_batch_idx] * - runtime_tokens_per_gen_step) + req_tokens_per_gen_step) + previous_batch_tokens_per_request.append( + req_tokens_per_gen_step) + extend_tokens_emitted += req_tokens_per_gen_step num_cached_tokens_per_seq.append( - past_seen_token_num + runtime_tokens_per_gen_step - + past_seen_token_num + req_tokens_per_gen_step - request.py_num_compressed_tokens) request.cached_tokens = (past_seen_token_num + - runtime_tokens_per_gen_step) + req_tokens_per_gen_step) if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( self.attn_backend) and spec_config.is_linear_tree: - prompt_lengths.append(runtime_tokens_per_gen_step) + prompt_lengths.append(req_tokens_per_gen_step) else: prompt_lengths.append(request.py_prompt_len) @@ -4479,6 +5264,15 @@ def append_cross_attention_state(request: LlmRequest, gen_request_seq_slots.append(request.py_seq_slot) previous_batch_len = len(previous_batch_indices) + # Device-window prologue precondition: it gathers this step's inputs + # by (slot, offset) through previous_batch_indices_cuda, which only + # covers requests that carried a previous device tensor and only in + # batch order when EVERY real generation request did. Recorded here, + # consumed by _apply_device_window_prologue. + self._dspark_prev_covers_batch = ( + previous_batch_len > 0 + and previous_batch_len == len(extend_requests) - + len(extend_dummy_requests)) def previous_seq_slots_device(): previous_batch_indices_host = torch.tensor( @@ -4628,30 +5422,79 @@ def previous_seq_slots_device(): if previous_batch_len > 0: previous_slots = previous_seq_slots_device() + # Ragged windows break the fixed strided gathers below; use + # explicit (row, col) index lists, keeping the cheaper strided + # path for uniform batches. + is_ragged_gen = any( + tokens != runtime_tokens_per_gen_step + for tokens in previous_batch_tokens_per_request) # previous input ids - previous_batch_tokens = (previous_batch_len * - runtime_tokens_per_gen_step) - new_tokens = new_tokens_device.transpose( - 0, - 1)[previous_slots, :runtime_tokens_per_gen_step].flatten() + previous_batch_tokens = len(previous_pos_indices) + if is_ragged_gen: + token_rows, token_cols = self._ragged_gather_indices( + previous_batch_indices, + previous_batch_tokens_per_request) + new_tokens = new_tokens_device.transpose( + 0, 1)[token_rows, token_cols].flatten() + else: + new_tokens = new_tokens_device.transpose(0, 1)[ + previous_slots, :runtime_tokens_per_gen_step].flatten() self.input_ids_cuda[num_tokens:num_tokens + previous_batch_tokens].copy_( new_tokens, non_blocking=True) # previous draft tokens - previous_batch_draft_tokens = (previous_batch_len * - runtime_draft_token_buffer_width) - if runtime_draft_token_buffer_width > 0: - self.draft_tokens_cuda[ - num_draft_tokens:num_draft_tokens + - previous_batch_draft_tokens].copy_( - next_draft_tokens_device[ - previous_slots, : - runtime_draft_token_buffer_width].flatten(), - non_blocking=True) - # prepare data for the preprocess inputs - kv_len_offsets_device = (new_tokens_lens_device - - runtime_tokens_per_gen_step) + if is_ragged_gen: + previous_batch_draft_tokens = (previous_batch_tokens - + previous_batch_len) + if previous_batch_draft_tokens > 0: + draft_rows, draft_cols = self._ragged_gather_indices( + previous_batch_indices, + [t - 1 for t in previous_batch_tokens_per_request]) + self.draft_tokens_cuda[ + num_draft_tokens:num_draft_tokens + + previous_batch_draft_tokens].copy_( + next_draft_tokens_device[draft_rows, + draft_cols].flatten(), + non_blocking=True) + else: + previous_batch_draft_tokens = ( + previous_batch_len * runtime_draft_token_buffer_width) + if runtime_draft_token_buffer_width > 0: + self.draft_tokens_cuda[ + num_draft_tokens:num_draft_tokens + + previous_batch_draft_tokens].copy_( + next_draft_tokens_device[ + previous_slots, : + runtime_draft_token_buffer_width].flatten(), + non_blocking=True) + # prepare data for the preprocess inputs. + # kv_len_offsets pairs with the num_cached_tokens_per_seq the + # host wrote above (past_seen + tokens_this_step): the two + # cancel to past_seen + accepted, so both sides must use the + # SAME per-request token count or the KV length is off by + # exactly their difference. + if is_ragged_gen: + tokens_per_previous_request = torch.tensor( + previous_batch_tokens_per_request, + dtype=torch.long, + pin_memory=prefer_pinned()).to( + new_tokens_lens_device.device, non_blocking=True) + previous_kv_len_offsets = ( + new_tokens_lens_device[previous_slots] - + tokens_per_previous_request) + # A slot the sampler has not yet written yields a stale + # count, and the captured kv_lens correction then walks the + # KV append out of bounds (the ragged IMA); clamp to the + # physically possible range. + previous_kv_len_offsets = previous_kv_len_offsets.clamp_( + min=-int(self.runtime_draft_len + 1), + max=int(self.runtime_draft_len + 1)) + else: + kv_len_offsets_device = (new_tokens_lens_device - + runtime_tokens_per_gen_step) + previous_kv_len_offsets = kv_len_offsets_device[ + previous_slots] previous_pos_indices_host = torch.tensor( previous_pos_indices, dtype=torch.int, @@ -4675,10 +5518,19 @@ def previous_seq_slots_device(): num_extend_reqeust_wo_dummy = len(extend_requests) - len( extend_dummy_requests) + if is_ragged_gen: + # previous_pos_id_offsets_cuda is indexed by flat + # generation-token position, which is no longer + # request_index * runtime_tokens_per_gen_step. Use the + # offset accumulated while walking the batch. + pos_offsets_start = previous_batch_token_start + else: + pos_offsets_start = ( + num_extend_reqeust_wo_dummy - + previous_batch_len) * runtime_tokens_per_gen_step self.previous_pos_id_offsets_cuda[ - (num_extend_reqeust_wo_dummy - previous_batch_len) * - runtime_tokens_per_gen_step:num_extend_reqeust_wo_dummy * - runtime_tokens_per_gen_step].copy_( + pos_offsets_start:pos_offsets_start + + previous_batch_tokens].copy_( new_tokens_lens_device[self.previous_pos_indices_cuda[ 0:previous_batch_tokens]], non_blocking=True) @@ -4686,8 +5538,7 @@ def previous_seq_slots_device(): self.previous_kv_lens_offsets_cuda[ num_extend_reqeust_wo_dummy - previous_batch_len:num_extend_reqeust_wo_dummy].copy_( - kv_len_offsets_device[previous_slots], - non_blocking=True) + previous_kv_len_offsets, non_blocking=True) elif new_tokens_device is not None: seq_slots_device = previous_seq_slots_device() @@ -4801,7 +5652,23 @@ def previous_seq_slots_device(): helix_is_inactive_rank=helix_is_inactive_rank, ) - if not attn_metadata.is_cuda_graph: + # Under ragged verification the per-request lengths move between + # replays of the same graph (the key pins the row count and token + # total, not the split), so seq_lens must be refreshed or + # attn_metadata.num_tokens disagrees with the input_ids width. Publish + # BEFORE the refresh decision below: it asks the metadata whether this + # step is ragged, and asking before this step's windows are on it + # reads the previous step's answer. + self._publish_gen_token_layout(attn_metadata, + scheduled_requests.generation_requests) + + refresh_seq_lens = not attn_metadata.is_cuda_graph + if (not refresh_seq_lens + and getattr(attn_metadata, "is_ragged_verify", False) + and attn_metadata.seq_lens is not None + and len(sequence_lengths) == attn_metadata.seq_lens.shape[0]): + refresh_seq_lens = True + if refresh_seq_lens: # Assumes seq lens do not change between CUDA graph invocations. This applies # to draft sequences too. This means that all draft sequences must be padded. attn_metadata.seq_lens = torch.tensor( @@ -4855,6 +5722,12 @@ def previous_seq_slots_device(): # pre-prepare counts so the steady-gen recording below stores values # that the per-step prepare() can re-clamp from scratch. num_cached_tokens_snapshot = list(num_cached_tokens_per_seq) + # The gen-token layout was published above, before the seq_lens refresh + # that depends on it. prepare() is its other consumer: it decides the DSA + # expanded-buffer layout, strides the expansions by this step's + # per-request token count, builds the per-row causal extents, and derives + # DeepSeek-V4's per-request compressor token counts. The rest of the + # layout is attached later, once the spec metadata exists. attn_metadata.prepare() cross_attention_inputs = (self._prepare_enc_dec_cross_attn_inputs( cross_encoder_hidden_states, @@ -4941,6 +5814,9 @@ def previous_seq_slots_device(): spec_metadata.seq_lens = sequence_lengths spec_metadata.num_accepted_draft_tokens = self.num_accepted_draft_tokens_cuda[:len( num_accepted_draft_tokens)] + self._attach_ragged_verify_layout( + spec_metadata, attn_metadata, + scheduled_requests.generation_requests) if isinstance(spec_metadata, Eagle3SpecMetadata): spec_metadata.request_accepted_path = request_accepted_path # No-op for non 1-model @@ -4983,6 +5859,10 @@ def previous_seq_slots_device(): if not self.is_warmup: self.previous_request_ids = all_gen_request_ids + self.previous_verify_lens = [ + getattr(request, "py_verify_len", None) + for request in scheduled_requests.generation_requests + ] self.has_previous_device_draft = next_draft_tokens_device is not None # Record the steady-state generation cache when this pass handled @@ -6204,6 +7084,15 @@ def forward(self, ) can_run_graph = key is not None + # A ragged step that misses its captured shape quietly runs eager; + # count it. Only generation-only steps count: a batch with context + # requests has no generation graph to miss. + if padded_requests.can_run_cuda_graph: + self._record_dspark_graph_use( + replayed=can_run_graph, + shape=None if can_run_graph else + (self.cuda_graph_runner.last_miss_reason or "unknown", + self.cuda_graph_runner.last_miss_key)) if can_run_graph: attn_metadata = maybe_attn_metadata spec_metadata = maybe_spec_metadata @@ -6231,6 +7120,16 @@ def forward(self, self._prepare_inputs_event = torch.cuda.Event() self._prepare_inputs_event.record() + # Device-window selection: re-rank the staged ragged layout by the + # verified block's own confidence, between input staging and the + # replay (stream-ordered, no host sync). Graph steps only: eager + # and capture steps run the host shape split, which is a valid + # window assignment. + if (can_run_graph and not self.is_warmup + and getattr(self, "_dspark_device_budget", None) + is not None): + self._apply_device_window_prologue(inputs, new_tensors_device) + with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): if not can_run_graph: # Fallback to eager execution if graph was not used diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 34ebd340a5ee..b1a058479d5a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2231,7 +2231,8 @@ def _append_iter_stats(self, attention_dp_rank: Optional[int] = None, host_step_time_ms: Optional[float] = None, prev_device_step_time_ms: Optional[float] = None, - gpu_forward_time_ms: Optional[float] = None): + gpu_forward_time_ms: Optional[float] = None, + num_generation_tokens: Optional[int] = None): """Append one iteration's finalized stats to the export buffer. The normal Attention-DP path fans out rank-local rows before calling @@ -2300,6 +2301,16 @@ def _append_iter_stats(self, local_dict["prevDeviceStepTimeMS"] = prev_device_step_time_ms if gpu_forward_time_ms is not None: local_dict["gpuForwardTimeMS"] = gpu_forward_time_ms + # The executed (padded) generation-token total -- the M axis of + # the SPS cost table. Same value the iteration log prints in + # `states`; carried per rank so ragged sweeps can align steps on + # the group's agreed shape. + num_gen_tokens = (num_generation_tokens + if num_generation_tokens is not None else + (self.model_engine.iter_states + or {}).get('num_generation_tokens')) + if num_gen_tokens is not None: + local_dict["numGenerationTokens"] = int(num_gen_tokens) local_dict["schedulerMode"] = scheduler_mode local_dict["rank"] = self.dist.tp_rank # Buffer for _flush_iter_stats_synced; in-place tp_allgather would @@ -2318,6 +2329,8 @@ def _append_iter_stats(self, # [5] prev_device_step_time_ms: Optional[float] # [6] scheduler_mode: "overlap" | "non_overlap" # [7] gpu_forward_time_ms: Optional[float] + # [8] num_generation_tokens: Optional[int] -- executed (padded) + # generation-token total, rank0's value broadcast per row with self.stats_lock: if (not _stats_buffer_is_unbounded(self.max_stats_len) and len(self.stats) > self.max_stats_len): @@ -2325,7 +2338,7 @@ def _append_iter_stats(self, self.stats.append( (stats, req_stats, kv_iter_stats, attention_dp_rank, host_step_time_ms, prev_device_step_time_ms, scheduler_mode, - gpu_forward_time_ms)) + gpu_forward_time_ms, num_generation_tokens)) def _process_iter_stats( self, @@ -2390,7 +2403,9 @@ def _process_iter_stats( is_rank0=self.dist.rank == 0, host_step_time_ms=host_step_time_ms, prev_device_step_time_ms=prev_device_step_time_ms, - gpu_forward_time_ms=gpu_forward_time_ms) + gpu_forward_time_ms=gpu_forward_time_ms, + num_generation_tokens=(self.model_engine.iter_states + or {}).get('num_generation_tokens')) else: self._append_iter_stats( stats, @@ -2399,7 +2414,28 @@ def _process_iter_stats( prev_device_step_time_ms=prev_device_step_time_ms, gpu_forward_time_ms=gpu_forward_time_ms) + def _report_dspark_ragged_stats(self) -> None: + """Final DSpark scheduling summary at shutdown. + + The counters live in the worker process and are unreachable from the + driver under multi-rank TP, so the log is the shared channel. + """ + worker = getattr(self.model_engine, "_get_spec_worker", None) + stats = getattr(worker(), "ragged_stats", None) if worker else None + # Flush BEFORE the ragged early-return: STS collection also runs in + # static mode. + recorder = getattr(self.sampler, "sts_recorder", None) + if recorder is not None: + recorder.flush() + if stats is None or not stats.steps_total: + return + stats.log_summary(prefix="DSpark ragged verify [final]") + def _executor_loop_cleanup(self): + try: + self._report_dspark_ragged_stats() + except Exception as exc: # noqa: BLE001 - never block shutdown on stats + logger.warning(f"DSpark ragged stats reporting failed: {exc}") # Wake any waiters in await_responses BEFORE potentially-blocking # work below. If wait_on_pp_send_handles hangs (e.g. after a # crash leaves PP send handles in a bad state), the await loop @@ -3186,6 +3222,401 @@ def wait_on_pp_send_handles(self, send_handles, microbatch_id): send_handles[microbatch_id].wait() send_handles[microbatch_id] = None + def set_dspark_verify_len_pin(self, verify_len: Optional[int]) -> Optional[int]: + """Queue a DSpark verify-length pin, or ``None`` to clear it. + + The pin bypasses the confidence planner so every request verifies the + same number of positions (for cost-table sweeps against a running + server). Validated here, applied on the next decode step, when the + step's cross-rank allgather hands the same value to every rank. + + Returns: + The validated pin, or None when cleared. + + Raises: + RuntimeError: if this executor has no DSpark confidence planner. + ValueError: if the length is outside the ladder (see + :meth:`DSparkVerifyPlanner.validate_verify_len_pin`). + """ + worker_fn = getattr(self.model_engine, "_get_spec_worker", None) + worker = worker_fn() if worker_fn else None + planner = getattr(worker, "verify_planner", None) if worker else None + if planner is None: + raise RuntimeError( + "no DSpark confidence planner on this executor; the verify-len " + "pin only applies to a DSpark engine with confidence " + "scheduling enabled") + mode = worker.ragged_verify_mode + computes = (mode.computes_windows if mode is not None else + self.model_engine.spec_config.enable_ragged_verify) + if not computes: + # Static mode never calls decide_verify_lens, so an accepted pin + # would be adopted and silently never applied. + raise RuntimeError( + "this engine runs static (uniform) verification; the " + "verify-length pin only affects window-computing modes, so " + "pinning here would be accepted and silently never applied") + if (worker.sts_recorder is not None + and verify_len is not None + and int(verify_len) < int(self.model_engine.max_draft_len)): + # A sub-block pin censors the acceptance label of every STS sample + # collected after it; the collector's start-time guard cannot see + # a runtime pin. + raise RuntimeError( + "STS calibration collection is active on this engine; a pin " + "below the full block censors the acceptance label for every " + "sample collected after it. Pin the full block, clear the " + "pin, or restart without the collector.") + if verify_len is not None and (planner.forced_budget_frac is not None + or planner.pending_budget_frac() >= 0): + raise RuntimeError( + "a budget fraction is in force (or queued) on this engine; a " + "pin would silently shadow it. Clear the fraction first.") + return planner.request_verify_len_pin(verify_len) + + def set_dspark_budget_frac(self, + frac: Optional[float]) -> Optional[float]: + """Queue a DSpark verify-budget fraction, or ``None`` to clear it. + + Replaces the cost-table argmax with ``frac`` of the maximum trimmable + budget while the confidence top-k keeps assigning the windows -- the + production ragged path at a controlled load point, for SPS cost-table + sweeps. Validated here, applied on the next decode step through the + step's cross-rank allgather. + + Raises: + RuntimeError: if this executor has no DSpark confidence planner, + runs static verification, or is collecting STS calibration + data (a trimmed budget censors the acceptance labels). + ValueError: if ``frac`` is outside ``(0.0, 1.0]``. + """ + worker_fn = getattr(self.model_engine, "_get_spec_worker", None) + worker = worker_fn() if worker_fn else None + planner = getattr(worker, "verify_planner", None) if worker else None + if planner is None: + raise RuntimeError( + "no DSpark confidence planner on this executor; the budget " + "fraction only applies to a DSpark engine with confidence " + "scheduling enabled") + mode = worker.ragged_verify_mode + trims = (mode.trims_submitted_tokens if mode is not None else + self.model_engine.spec_config.enable_ragged_verify) + if not trims: + # Static never computes windows; cap-accept computes them but + # submits the full block, so a frac there cannot move the load + # point -- it would only cap live acceptance, which is not what a + # cost sweep asks for. + raise RuntimeError( + "this engine does not trim submitted tokens (static or " + "cap-accept mode); the budget fraction only moves the load " + "point in compact mode, so setting it here would change " + "nothing or change acceptance instead") + if frac is not None and (planner.forced_verify_len is not None + or planner.pending_verify_len_pin() >= 0): + raise RuntimeError( + "a verify-length pin is in force (or queued) on this engine; " + "the pin takes precedence and the fraction would be adopted, " + "logged, and never applied. Clear the pin first.") + if (worker.sts_recorder is not None and frac is not None + and float(frac) < 1.0): + raise RuntimeError( + "STS calibration collection is active on this engine; a " + "trimmed budget censors the acceptance label for every sample " + "collected after it. Use 1.0, clear the fraction, or restart " + "without the collector.") + return planner.request_budget_frac(frac) + + def _dspark_confidence_draft_len(self, + scheduled_batch: ScheduledRequests) -> int: + """Draft length for this iteration, from the DSpark confidence planner. + + Runs outside CUDA-graph capture: it reads the confidence snapshot + staged on the previous iteration, then stages a fresh one -- a + one-iteration lag with no synchronization. Falls back to the full + block on any missing piece. + """ + max_draft_len = self.model_engine.max_draft_len + worker = self.model_engine._get_spec_worker() + planner = getattr(worker, "verify_planner", None) + if planner is None: + return max_draft_len + + gen_requests = scheduled_batch.generation_requests + num_gen_requests = len(gen_requests) + # Read back by SLOT, not batch position: the snapshot is one iteration + # old and the batch reshuffles in between. Slotless requests map to + # the neutral row ("verify the whole block"). + confidence_rows = [ + worker.confidence_row_for(request.py_request_id) + for request in gen_requests + ] + + # The mode, not the config flag, decides whether windows are computed; + # every counter below is keyed off the same decision. + stats = worker.ragged_stats + # Wire the worker's counters into the sampler here, where both are + # reachable: acceptance is only observable host-side in the sampler. + if (stats is not None and self.sampler is not None + and self.sampler.acceptance_stats is None): + self.sampler.acceptance_stats = stats + # Same wiring for STS. The sampler gets the worker's OWN row resolver + # and pass counter: the recorder joins by (draft pass, buffer row), + # and both keys must come from the allocator that wrote the buffer. + if (self.sampler is not None + and self.sampler.sts_recorder is None): + recorder = worker.sts_recorder + if recorder is not None: + self.sampler.sts_recorder = recorder + self.sampler.sts_row_for = worker.sts_row_for + self.sampler.sts_seq_provider = worker.verified_draft_seq + mode = worker.ragged_verify_mode + compute_windows = mode.computes_windows + + is_distributed = (self.dist is not None + and self.dist.tp_size > 1) + + # --- Phase 1: decide locally, issue nothing --------------------------- + # + # Every decline reason is rank-local, so a collective placed after + # them deadlocks: nothing below may branch on rank-local state before + # the single collective, and nothing after it may issue another one. + ragged_lens = None + fallback_reason = None + device_budget = None + # Cleared unconditionally: a budget left over from a previous step + # would re-rank a batch it was never computed for. + self.model_engine._dspark_device_budget = None + if compute_windows: + if planner.skip_mixed_trim and scheduled_batch.context_requests: + # Mixed-step guard (experiment): prefill work shares this + # step, so the pure-decode cost table cannot price the trim. + # Declining here rides the normal all-or-nothing gate below. + fallback_reason = "mixed_step_declined" + elif planner.device_windows and mode.trims_submitted_tokens: + # Device-window mode: the host decides only the CAPACITY, from + # the lag-2 snapshot; the returned lens are a shape split used + # for the cross-rank agreement and the fit. The true windows + # are re-ranked on device just before the replay. + pair = planner.decide_verify_budget( + num_gen_requests=num_gen_requests, rows=confidence_rows) + if pair is None: + fallback_reason = "planner_declined" + else: + device_budget, ragged_lens = pair + elif planner.device_windows: + # cap-accept + device mode: a pure measurement configuration. + # The lag-2 argmax picks the budget, the normal top-k spends + # it, and the block still executes IN FULL -- the windows ride + # py_verify_cap as accounting caps, so cap_trim_tokens reports + # exactly how many target-accepted tokens this budget would + # have discarded. Differencing against plain cap-accept + # (lag-1 budget) isolates the budget's true acceptance cost + # at identical compute. + pair = planner.decide_verify_budget( + num_gen_requests=num_gen_requests, rows=confidence_rows) + ragged_lens = (planner.decide_verify_lens( + num_gen_requests=num_gen_requests, + rows=confidence_rows, + reduce_across_ranks=False, + budget_override=pair[0]) if pair is not None else None) + if ragged_lens is None: + fallback_reason = "planner_declined" + else: + ragged_lens = planner.decide_verify_lens( + num_gen_requests=num_gen_requests, + rows=confidence_rows, + reduce_across_ranks=False) + if ragged_lens is None: + fallback_reason = "planner_declined" + else: + fallback_reason = "mode_static" + # Nothing to compute: the fallback verifies the full drafted block. + + # --- Phase 2: one collective, same shape on every rank, every step ---- + # + # Payload: [wants_ragged, num_rows, total_verify_tokens, max_window, + # can_graph, pending_pin]. `can_graph` rides along because the ragged + # fit runs before the graph runner's own allgather and must know + # whether row padding will happen: _get_padded_batch only takes the + # cross-rank row maximum when EVERY rank can graph this step. + local_rows = len(ragged_lens) if ragged_lens is not None else 0 + local_tokens = (sum(1 + int(v) + for v in ragged_lens) if ragged_lens else 0) + local_max_window = max(ragged_lens) if ragged_lens else 0 + peer_stats = None + # Single-rank runs have no peer that could be holding a context + # request, so the local batch's own answer is the group's. + all_can_graph = bool(scheduled_batch.can_run_cuda_graph) + if is_distributed: + payloads = self.dist.tp_allgather([ + 1 if ragged_lens is not None else 0, + local_rows, + local_tokens, + local_max_window, + 1 if scheduled_batch.can_run_cuda_graph else 0, + planner.pending_verify_len_pin(), + planner.pending_budget_frac(), + ]) + # The pin and the budget fraction ride the allgather so every rank + # adopts the same value on the same step (the endpoint lands on + # one rank). -1 means nobody queued anything. + agreed_pin = next( + (int(payload[5]) for payload in payloads + if len(payload) > 5 and int(payload[5]) >= 0), -1) + planner.adopt_verify_len_pin(agreed_pin) + agreed_frac = next( + (int(payload[6]) for payload in payloads + if len(payload) > 6 and int(payload[6]) >= 0), -1) + planner.adopt_budget_frac(agreed_frac) + if not all(int(payload[0]) for payload in payloads): + # All-or-nothing across ranks: `all_rank_num_tokens` is frozen + # into the captured graph, so a mixed step's mismatch would + # not even be observable at replay time. + if ragged_lens is not None: + fallback_reason = "peer_declined" + ragged_lens = None + else: + # payload[3] is unread; it stays so the collective keeps a + # stable shape. peer_stats' third element is the GROUP's + # can-graph answer, identical on every rank, not each rank's + # own flag. + all_can_graph = all(int(payload[4]) for payload in payloads) + peer_stats = [[ + int(payload[1]), + int(payload[2]), 1 if all_can_graph else 0 + ] for payload in payloads] + else: + # Single-rank world: the local queue IS the group's agreement. + planner.adopt_verify_len_pin(planner.pending_verify_len_pin()) + planner.adopt_budget_frac(planner.pending_budget_frac()) + + # --- Phase 3: no more collectives ------------------------------------- + # + # Land the token total on a captured bucket BEFORE the graph key is + # built; no captured shape means fall through to uniform, which always + # has a graph. The fit can still differ across ranks, which is why the + # attention-DP graph gate compares the token bucket, not just the + # batch size. + # + # `cap-accept` submits the full block: its windows go on + # `py_verify_cap`, not `py_verify_len`, so the whole token-layout path + # downstream keeps seeing an untrimmed batch. + bucket = None + ragged_active = False + cap_accept = (mode is not None and mode.computes_windows + and not mode.trims_submitted_tokens) + if cap_accept: + # Clear every request first, set second: a stale cap silently + # commits against last step's schedule, with no assert to trip. + for request in gen_requests: + request.py_verify_cap = None + if ragged_lens is not None: + if cap_accept: + for request, window in zip(gen_requests, ragged_lens): + request.py_verify_cap = int(window) + ragged_active = True + elif all_can_graph: + bucket = self.model_engine.fit_ragged_verify_lens( + gen_requests, ragged_lens, peer_stats=peer_stats) + ragged_active = bucket is not None + if bucket is None: + fallback_reason = "no_captured_shape" + elif device_budget is not None: + # The fit landed on a captured shape; hand the engine the + # capacity so the pre-replay prologue re-ranks it with + # fresh confidence. Published only here: every fallback + # leaves it None and the step runs the host windows. + self.model_engine._dspark_device_budget = int( + device_budget) + else: + # No rank can graph this step, so run the windows eagerly. Do + # not fit a bucket: the fit sizes pad rows _get_padded_batch + # will not append on an ungraphable step. + runner = self.model_engine.cuda_graph_runner + runner.agreed_ragged_bucket = None + runner.ragged_pad_verify_len = 0 + for request, window in zip(gen_requests, ragged_lens): + request.py_verify_len = int(window) + ragged_active = True + + # Record what the step actually decided. + if stats is not None: + # Periodic summary: under TP the counters sit across an MPI + # boundary, so the log is the only shared channel. + if stats.steps_total and stats.steps_total % 32 == 0: + stats.log_summary(prefix="DSpark ragged verify [periodic]") + fitted = None + delivered = None + if cap_accept: + # Nothing was fitted: the scheduled windows ARE the decision, + # and `delivered` must say the block was submitted whole. + fitted = [int(v) for v in ragged_lens] if ragged_lens else None + delivered = num_gen_requests * (1 + max_draft_len) + elif ragged_active: + # Read back off the requests, not `ragged_lens`: the fit spent + # the bucket's rounding slack on real verification. + fitted = [ + int(request.py_verify_len) for request in gen_requests + if getattr(request, "py_verify_len", None) is not None + ] + if len(fitted) != num_gen_requests: + fitted = None + stats.record_step(num_gen_requests=num_gen_requests, + verify_lens=fitted, + bucket=bucket, + # Only when the fit delivered: otherwise this + # books the LAST fit's row count on steps that + # never entered the fit. + padded_bs=(self.model_engine._dspark_last_padded_bs + if bucket is not None else None), + fallback=fallback_reason, + delivered=delivered) + + if ragged_active and not cap_accept: + # The block is always drafted in full, so the batch-wide draft + # length stays at the top tier (fixed-width buffers); per-request + # windows ride on py_verify_len. + runtime_draft_len = max( + int(t) for t in self.model_engine.spec_config.verify_len_tiers) + else: + # Clear stale windows on EVERY generation request: a batch where + # only some requests carry one gets a ragged input layout with + # uniform acceptance -- silent token misattribution. + for request in gen_requests: + request.py_verify_len = None + # Clear the published fit state too: a stale agreed bucket + # describes a batch that no longer exists. + if self.model_engine.cuda_graph_runner is not None: + self.model_engine.cuda_graph_runner.agreed_ragged_bucket = None + self.model_engine.cuda_graph_runner.ragged_pad_verify_len = 0 + # The only fallback is the full drafted block; nothing to agree on + # across ranks, so no collective. + runtime_draft_len = planner.max_tier + + # Stage this step's confidence for the *next* decision. Non-blocking; if + # it has not landed by then the planner just verifies the full block. + # The whole slot-indexed buffer is staged, so the next step can look up + # whichever requests it happens to be scheduling. + buffer = worker.staged_confidence_buffer() + if buffer is not None: + # The stamp buffer says which draft pass last wrote each slot; the + # sequence bump marks "everything drafted so far" so the planner + # can histogram how stale each gathered row was at decision time. + staged_seq = worker.bump_draft_seq() + stamps = worker.confidence_stamp_buffer() + planner.stage_confidence(buffer, stamps=stamps, + staged_seq=staged_seq) + # The STS recorder rings its own copy: labels arrive steps later + # and must pair with the pass that DRAFTED the block, not whatever + # the live buffer holds by then. + recorder = worker.sts_recorder + if recorder is not None: + recorder.stage_snapshot(device_logits=buffer, + device_stamps=stamps, + staged_seq=staged_seq) + + return int(runtime_draft_len) + def _handle_dynamic_draft_len(self, scheduled_batch: ScheduledRequests) -> None: """Handle dynamic draft length for the current batch. @@ -3213,19 +3644,31 @@ def _handle_dynamic_draft_len(self, self.model_engine.runtime_draft_len = 0 return - if (self.model_engine.spec_config is not None - and self.model_engine.spec_config.draft_len_schedule is not None - and self.model_engine.spec_config.spec_dec_mode. - support_dynamic_draft_len()): + spec_config = self.model_engine.spec_config + schedule_driven = ( + spec_config is not None + and spec_config.draft_len_schedule is not None + and spec_config.spec_dec_mode.support_dynamic_draft_len()) + # DSpark drives the draft length from the batch's confidence rather than + # from a static batch-size schedule, but consumes the exact same + # pad/truncate machinery below. + confidence_driven = (spec_config is not None and getattr( + spec_config, "enable_confidence_scheduling", False)) + + if schedule_driven or confidence_driven: from tensorrt_llm._torch.speculative.utils import \ get_draft_len_for_batch_size - spec_dec_mode = self.model_engine.spec_config.spec_dec_mode + spec_dec_mode = spec_config.spec_dec_mode - # 1. Resolve runtime draft length from schedule - runtime_draft_len = get_draft_len_for_batch_size( - self.model_engine.spec_config.draft_len_schedule, - scheduled_batch.batch_size, self.model_engine.max_draft_len) + # 1. Resolve runtime draft length + if confidence_driven: + runtime_draft_len = self._dspark_confidence_draft_len( + scheduled_batch) + else: + runtime_draft_len = get_draft_len_for_batch_size( + spec_config.draft_len_schedule, scheduled_batch.batch_size, + self.model_engine.max_draft_len) # 2. Pad or truncate draft tokens to the resolved length DRAFT_BUFFER_PAD = 0 # Buffer sentinel, not PARD mask_token_id. rejection_on = getattr(self.model_engine.spec_config, @@ -5024,7 +5467,8 @@ def _fetch_new_requests( host_step_time_ms=record.host_step_time_ms, prev_device_step_time_ms=record. prev_device_step_time_ms, - gpu_forward_time_ms=record.gpu_forward_time_ms) + gpu_forward_time_ms=record.gpu_forward_time_ms, + num_generation_tokens=record.num_generation_tokens) all_ranks_num_active_requests = [ s.num_active_requests for s in all_rank_states ] diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 3f649ca0db83..c221380490db 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -800,7 +800,10 @@ def add_token( def int_tensor(shape: tuple[int, ...], device: str = "cuda") -> torch.Tensor: - return torch.empty(shape, dtype=torch.int, device=device) + # zeros, not empty: these buffers are slot-indexed stores whose consumers + # (e.g. the overlap scheduler's previous-batch gather) may read slots the + # current step never wrote; an unwritten slot must read as 0. + return torch.zeros(shape, dtype=torch.int, device=device) @dataclass(kw_only=True, frozen=True) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 7f43f103ec4e..6d1ba60c3f7b 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -13,11 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # -# DSpark worker / metadata mirror the DFlash plumbing (capture target-layer -# hidden states, accept the previous block with standard verification, draft a -# new block in one backbone forward), adapted to DSpark's draft model which -# produces the whole block (and its confidence-truncated length) inside a single -# ``DSparkDraftModel.forward`` rather than via mask-token cross-attention. +# DSpark speculative decoding: capture target-layer hidden states during the +# target forward, verify the previous block, and draft a new block in a single +# ``DSparkDraftModel.forward``. from collections import deque from dataclasses import dataclass @@ -30,26 +28,25 @@ from tensorrt_llm.mapping import Mapping from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID +from .dspark_schedule import NEUTRAL_CONFIDENCE_LOGIT from .interface import SpecMetadata, SpecWorkerBase if TYPE_CHECKING: from ...llmapi.llm_args import DSparkDecodingConfig +# Unscored-slot fill: sigmoid saturates to 1.0, so an unwritten row schedules +# as "verify the whole block" (fail-safe). +_NEUTRAL_CONFIDENCE_LOGIT = NEUTRAL_CONFIDENCE_LOGIT + @dataclass class DSparkSpecMetadata(SpecMetadata): """Metadata for DSpark speculative decoding. - Captures hidden states from the target model's ``layers_to_capture`` during - the target forward pass. DSpark captures the *mean over the multi-head - (mHC) residual streams* at each captured layer (handled by the target-side - capture hook), concatenated across layers, and feeds them to the draft - model's ``main_proj`` + ``main_norm`` (inside ``DSparkDraftModel.forward``) - as the captured-context attention input (``main_x``). - - Mirrors :class:`DFlashSpecMetadata`; the only DSpark-specific detail is that - the per-layer captured width is the model hidden size (post hc-mean), so the - buffer is ``[max_num_tokens, hidden_size * num_capture_layers]``. + Captures target-layer hidden states (mean over the mHC residual streams, + per-layer width = model hidden size) into a + ``[max_num_tokens, hidden_size * num_capture_layers]`` buffer that feeds + ``DSparkDraftModel.forward``. """ batch_indices_cuda: Optional[torch.Tensor] = None @@ -71,11 +68,9 @@ def __post_init__(self): self.is_spec_dec_tree = False self.is_spec_dec_dynamic_tree = False - # Set up hidden state capture buffer if self.layers_to_capture is not None and len(self.layers_to_capture) > 0: self.layers_to_capture = sorted(list(self.layers_to_capture)) self.num_capture_layers = len(self.layers_to_capture) - # O(1) lookups for is_layer_capture() and maybe_capture_hidden_states() self._capture_layer_set = frozenset(self.layers_to_capture) self._layer_to_idx = {lid: i for i, lid in enumerate(self.layers_to_capture)} self.captured_hidden_states = torch.empty( @@ -100,10 +95,9 @@ def prepare(self): ) self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) - # CUDA-graph-safe path: maintain the request->slot mapping on the host - # (outside the captured region) and mirror it into ``_batch_to_slot`` so the - # captured gen forward can index the rolling windows by tensor. Mirrors - # ``DFlashSpecMetadata.prepare`` (dflash.py:96-113). + # Maintain the request->slot map on the host (outside any captured + # region) and mirror it into ``_batch_to_slot`` so the captured gen + # forward indexes the rolling windows by tensor. worker = getattr(self, "_dspark_worker", None) if worker is not None and worker._win_inited: current = set(self.request_ids) @@ -113,16 +107,11 @@ def prepare(self): worker._ctx_len[slot] = 0 worker._kv_windows[slot].zero_() worker._free_slots.append(slot) - # Assign a persistent rolling-window slot to every real generation - # request that never ran a context/seed forward on this worker. In - # disaggregated serving the prompt is prefilled (and the window - # seeded) on the *context* server, so ``_seed_context_windows`` never - # runs on the generation server and ``_req_to_slot`` stays empty; - # without this, all concurrent gen requests fall through to the shared - # scratch row below and corrupt each other's draft window at batch - # size > 1 (GitHub #16767). Context-prefix entries are left to - # ``_seed_context_windows``; the ADP-idle (id 0) and CUDA-graph - # padding dummies are kept on the scratch row. + # Assign a slot to every real gen request that never ran a + # context/seed forward on this worker (disaggregated serving + # prefills on the context server), so concurrent gen requests do + # not share the scratch row. Context requests are seeded by + # ``_seed_context_windows``; dummies stay on the scratch row. num_contexts = max(0, len(self.request_ids) - self.num_generations) for rid in self.request_ids[num_contexts:]: if ( @@ -131,10 +120,9 @@ def prepare(self): and rid not in worker._req_to_slot ): worker._assign_slot(rid, reset=False) - # Unknown request IDs (synthetic warmup / CUDA-graph padding, ADP idle - # requests, or disagg seed forwards without a real id) map to the - # dedicated throwaway scratch row so they cannot overwrite a live - # request's rolling window (they previously aliased to slot 0). + # Unknown request IDs (warmup / CUDA-graph padding, ADP idle, disagg + # seeds without a real id) map to the throwaway scratch row so they + # cannot overwrite a live request's rolling window. scratch = worker._scratch_slot mapping = torch.tensor( [worker._req_to_slot.get(rid, scratch) for rid in self.request_ids], @@ -150,14 +138,11 @@ def is_layer_capture(self, layer_id: int) -> bool: def maybe_capture_hidden_states( self, layer_id: int, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None ) -> None: - """Capture hidden states from a target model layer into the buffer. - - DeepSeek-V4 keeps the multi-head (mHC) residual stream flattened as - ``[num_tokens, hc_mult * hidden]``; DSpark captures the *mean over the hc - streams* (reference ``h.mean(dim=2)`` with ``h`` shaped - ``[*, hc_mult, hidden]``). We reduce here so the V4 decoder layer's - existing capture call is unchanged. A ``[num_tokens, hidden]`` input - (already reduced / non-mHC) is stored as-is. + """Capture a target layer's hidden states into the buffer. + + An mHC-flattened ``[num_tokens, hc_mult * hidden]`` input is reduced by + mean over the hc streams; a ``[num_tokens, hidden]`` input is stored + as-is. """ if self.captured_hidden_states is None: return @@ -165,7 +150,6 @@ def maybe_capture_hidden_states( if i is not None: num_tokens = hidden_states.shape[0] to_save = hidden_states + residual if residual is not None else hidden_states - # mHC residual -> mean over the hc_mult streams. if to_save.shape[-1] != self.hidden_size: hc_mult = to_save.shape[-1] // self.hidden_size to_save = to_save.reshape(num_tokens, hc_mult, self.hidden_size).mean(dim=1) @@ -185,27 +169,12 @@ def get_hidden_states(self, num_tokens: int) -> Optional[torch.Tensor]: class DSparkWorker(SpecWorkerBase): """Worker for DSpark speculative decoding. - DSpark drafts a whole block of ``block_size`` tokens in one backbone forward - (``DSparkDraftModel.forward``): it projects the captured target-layer hidden - states (``main_proj`` + ``main_norm``) into the draft's captured-context - attention, runs the ``num_stages`` DSpark blocks over a rolling captured - window, refines the per-position logits with the Markov head, and predicts a - per-position acceptance confidence used to truncate the proposed prefix. - - Unlike DFlash, the draft does NOT use the paged KV cache or mask-token - cross-attention: its attention K/V come from the worker-owned rolling window - of projected captured context (one ``main_kv`` per decode step, per stage). - Acceptance of the previous block goes through the unified - :meth:`SpecWorkerBase.sample_and_accept_draft_tokens` (strict target-verify, - or rejection sampling for a non-greedy batch), so greedy parity with no-spec - is preserved regardless of draft quality. - - The rolling window is kept consistent across the whole decode: it is seeded - from the prompt's captured context at prefill and back-filled with the - intermediate accepted tokens of a multi-accept step (both via - ``DSparkDraftModel.write_context_windows``), in addition to the per-step bonus - write done by the generation path. These affect draft acceptance rate only, - not correctness, which the standard target verify guarantees. + Drafts a whole block of ``block_size`` tokens in one + ``DSparkDraftModel.forward`` over a worker-owned rolling window of + projected captured context (no paged KV cache). Acceptance goes through + :meth:`SpecWorkerBase.sample_and_accept_draft_tokens`, so the window state + affects acceptance rate only, never correctness. Confidence scores feed + only the verification scheduler's budget, never the acceptance decision. Reference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec). """ @@ -227,31 +196,130 @@ def __init__( self._ctx_len: Optional[torch.Tensor] = None # [max_batch] abs decode position self._win = 0 - # Slot management. ``_req_to_slot`` (python dict) + ``_free_slots`` are the - # source of truth, updated in prepare()/forward(); ``_batch_to_slot`` is the - # CUDA mirror (request-order -> slot) read by the CUDA-graph-safe batched - # gen path (set on the host in prepare(), so the captured forward indexes - # the rolling windows through a tensor instead of a python dict lookup). + # Slot management: ``_req_to_slot`` + ``_free_slots`` (host) are the + # source of truth; ``_batch_to_slot`` is the CUDA mirror + # (request-order -> slot) the captured gen forward indexes through. self._req_to_slot = {} # request_id -> slot index self._free_slots = deque() # available slot indices self._batch_to_slot: Optional[torch.Tensor] = None # [max_batch] long, cuda - # Index of the throwaway "scratch" window row that absorbs padded / - # unknown request IDs (set in ``_lazy_init`` to ``max_batch``); it is - # never handed out through ``_free_slots``. + # Throwaway window row for padded / unknown request IDs (set to + # ``max_batch`` in ``_lazy_init``); never handed out via ``_free_slots``. self._scratch_slot = 0 - # The generation draft path is the batched, host-sync-free - # ``_draft_gen_block_batched`` + ``DSparkDraftModel.forward_batched`` + - # ``dspark_attention_forward_batched``: it is correct in eager mode AND safe - # to capture into the target's CUDA graph (DSpark is a one-engine drafter — - # its worker forward runs inside that graph, so the draft path MUST be - # capture-safe whenever ``cuda_graph_config`` is set). + # ``return_confidence`` is read once here (never per step) so the + # captured draft graph cannot diverge. ``_confidence_logits`` is the + # slot-indexed handoff to the verification scheduler. + self.return_confidence = bool(spec_config.enable_confidence_scheduling) + import os as _os + + from .dspark_sts import STS_COLLECT_ENV + if _os.environ.get(STS_COLLECT_ENV) and not self.return_confidence: + raise ValueError( + f"{STS_COLLECT_ENV} is set but enable_confidence_scheduling is " + f"off; the recorder is only built on the confidence path, so " + f"this run would complete cleanly and collect NOTHING. Enable " + f"confidence scheduling (a missing cost table keeps the " + f"planner on the full block) or unset the variable.") + self._confidence_logits: Optional[torch.Tensor] = None # [max_batch+2, block] + # Draft-pass stamps for the buffer above; allocated in ``_lazy_init``. + self._confidence_stamp: Optional[torch.Tensor] = None + self._draft_seq_host = 0 + self._draft_seq_cuda: Optional[torch.Tensor] = None + # Row index of the permanently-neutral confidence row; see ``_lazy_init``. + self._neutral_conf_row = 0 + # Host-side planner turning the confidence snapshot into this + # iteration's draft length; built in ``_lazy_init``. + self.verify_planner = None + + # DSpark is a one-engine drafter: the worker forward runs inside the + # target's CUDA graph, so the gen draft path + # (``_draft_gen_block_batched``) must stay host-sync-free and + # capture-safe whenever ``cuda_graph_config`` is set. logger.info( f"DSparkWorker initialized with " f"use_separate_draft_kv_cache={use_separate_draft_kv_cache}" ) + def staged_confidence_buffer(self) -> Optional[torch.Tensor]: + """The whole slot-indexed confidence buffer, or None when disabled. + + The whole buffer, not the current batch's rows: the planner reads it + back by slot across the one-iteration lag. + """ + return self._confidence_logits + + def confidence_stamp_buffer(self) -> Optional[torch.Tensor]: + """The slot-indexed draft-pass stamps, or None before ``_lazy_init``.""" + return self._confidence_stamp + + def bump_draft_seq(self) -> int: + """Advance the draft-pass sequence and return the PREVIOUS value. + + Must be called once per step from the executor's host path before the + step's forward is enqueued; the returned previous value is the + sequence of the last completed draft. + """ + prev = self._draft_seq_host + self._draft_seq_host += 1 + if self._draft_seq_cuda is not None: + self._draft_seq_cuda.fill_(self._draft_seq_host) + return prev + + def confidence_row_for(self, req_id: int) -> int: + """Buffer row holding ``req_id``'s confidence, host-side. + + Falls back to the permanently-neutral row ("verify the full block" -- + fail-safe). Keyed by slot, not batch position, so the one-iteration- + lagged snapshot survives batch reshuffles. + """ + if self._confidence_logits is None: + return self._neutral_conf_row + return self._req_to_slot.get(req_id, self._neutral_conf_row) + + def sts_row_for(self, req_id: int) -> Optional[int]: + """``confidence_row_for`` for the STS recorder: None, never neutral. + + Pairing a label against the never-written neutral row would fabricate + a sample; an unresolvable request is dropped (counted as ``no_row``). + """ + if self._confidence_logits is None: + return None + return self._req_to_slot.get(req_id) + + def verified_draft_seq(self) -> Optional[int]: + """The draft pass that produced the block the CURRENT step verifies. + + Read at sampling time, after this step's ``bump_draft_seq``, so the + verified pass is the previous value. Passes start at 1: a sequence of + 0 would false-match every never-written stamp row. + """ + return self._draft_seq_host - 1 if self._draft_seq_host > 1 else None + + def verified_draft_seq_cuda(self) -> Optional[torch.Tensor]: + """:meth:`verified_draft_seq` as a device scalar, for the prologue. + + ``_draft_seq_cuda`` already holds the CURRENT step's sequence by + launch time (``bump_draft_seq`` fills it before the forward is + enqueued), so the pass that stamped the block being verified is one + less. Passing the buffer itself would make every fresh row read as + stale and silently degrade the whole batch to neutral. + """ + if self._draft_seq_cuda is None: + return None + return self._draft_seq_cuda - 1 + + def batch_slot_view(self, num_rows: int) -> Optional[torch.Tensor]: + """The batch-position -> confidence-row map staged this step. + + ``[num_rows]`` int64 on device; rows past the real batch point at the + scratch row (their survival is zeroed by the prologue's ``num_real`` + mask, so the content never matters). + """ + if self._batch_to_slot is None: + return None + return self._batch_to_slot[:num_rows] + @property def max_draft_len(self) -> int: return self.spec_config.max_draft_len @@ -271,23 +339,18 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: self._win = int(draft_model._attn_params["window_size"]) head_dim = int(draft_model._attn_params["head_dim"]) - # Real requests occupy slots ``[0, max_batch)``; one extra "scratch" row - # at index ``max_batch`` absorbs padded / unknown request IDs (CUDA-graph - # padding, ADP idle requests, or disagg seed forwards that arrive without - # a real request id) so they can never overwrite a live request's rolling - # window. Previously such IDs aliased to slot 0 and corrupted whichever - # real request occupied it. The scratch row is never handed out through - # ``_free_slots`` and its contents are throwaway. + # Real requests occupy slots ``[0, max_batch)``; the extra scratch row + # at index ``max_batch`` absorbs padded / unknown request IDs so they + # can never overwrite a live request's rolling window. It is never + # handed out through ``_free_slots``. self._scratch_slot = max_batch num_rows = max_batch + 1 - # CUDA-graph padding requests carry ids in - # ``[CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, CUDA_GRAPH_DUMMY_REQUEST_ID]``, - # while real request ids start at ``max_batch_size`` and grow, so a simple - # floor cleanly separates them. Together with ``ATTENTION_DP_DUMMY_REQUEST_ID`` - # (0) these dummies must route to the scratch row (see ``prepare()``) and - # never consume a real slot. Imported lazily to break the - # dspark -> cuda_graph_runner -> speculative.utils -> dspark import cycle. + # CUDA-graph padding ids sit in + # ``[CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, CUDA_GRAPH_DUMMY_REQUEST_ID]``; + # real ids start at ``max_batch_size``, so a floor separates them. + # Imported lazily to break the dspark -> cuda_graph_runner -> + # speculative.utils -> dspark import cycle. from ..pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID self._graph_dummy_id_floor = CUDA_GRAPH_DUMMY_REQUEST_ID - self.max_draft_len @@ -299,6 +362,34 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: ) self._ctx_len = torch.zeros(num_rows, dtype=torch.long, device="cuda") self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") + if self.return_confidence: + # Neutral fill = large positive logit (sigmoid ~ 1.0): unscored + # slots schedule as "verify the full block" -- the fail-safe + # direction. The neutral row sits one past the scratch row; it is + # never in ``_batch_to_slot``, so no draft ever writes it. + self._neutral_conf_row = num_rows + self._confidence_logits = torch.full( + (num_rows + 1, block_size), + _NEUTRAL_CONFIDENCE_LOGIT, + dtype=torch.float32, + device="cuda", + ) + # Per-row freshness stamps: the draft-pass sequence of the last + # write, distinguishing fresh data from a stale replay or the + # neutral fill. Written in-graph through the same ``slots`` scatter + # as the confidence; the scalar is bumped outside the graph. + self._confidence_stamp = torch.zeros( + (num_rows + 1,), dtype=torch.int32, device="cuda") + self._draft_seq_host = 0 + # Bumped by ``fill_`` from the executor's host path, outside + # inference mode -- an inference tensor refuses that update, so + # allocate this scalar with inference mode off. The stamp buffer + # above intentionally stays an inference tensor (written in-graph + # only). + with torch.inference_mode(False): + self._draft_seq_cuda = torch.zeros((), dtype=torch.int32, + device="cuda") + self._build_verify_planner(draft_model, block_size, max_batch) self._free_slots = deque(range(max_batch)) self._req_to_slot = {} self._win_inited = True @@ -308,6 +399,158 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: f"({max_batch} request slots + 1 scratch row)" ) + def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> None: + """Construct the host-side verify planner (once, at lazy init). + + ``max_batch`` is only used to derive the verify-length tier ladder when + the user did not configure one; see :func:`derive_verify_len_tiers` for + why the ladder is a function of batch size. + """ + import json + + from .dspark_planner import SpsCostTable + from .dspark_schedule import DSparkScheduleConfig + from .dspark_verify import DSparkVerifyPlanner + + cfg = DSparkScheduleConfig(block_size=block_size, min_verify_len=1) + + cost_table = None + table_path = self.spec_config.confidence_sps_table_path + if table_path: + with open(table_path, encoding="utf-8") as f: + raw = json.load(f) + # The overhead fields are optional but not cosmetic: omitting them + # asserts ``step_time_ms`` already measures whole steps at the + # deployment's batch size. + from .dspark_planner import check_table_fingerprint + if (raw.get("_meta") or {}).get("lookup") != "interp": + logger.warning( + f"DSpark cost table {table_path} carries no " + f"lookup='interp' marker: it may predate the " + f"interpolating consumer. Tables written for the old " + f"floor lookup dropped shelf-closing breakpoints on " + f"purpose, and interpolating across those gaps re-bills " + f"every mid-shelf total upward. Re-emit the table with a " + f"current profiler if trimming behaves oddly.") + check_table_fingerprint(payload=raw, live={ + "tp": int(self.mapping.tp_size), + "ep": int(self.mapping.moe_ep_size), + "attention_dp": bool(self.mapping.enable_attention_dp), + "block": int(block_size), + "max_batch_size": int(max_batch), + # moe_backend deliberately absent: a wrong live value would + # false-reject a correct table; the table's value surfaces on + # the "check manually" INFO line instead. + }) + cost_table = SpsCostTable( + token_counts=[int(v) for v in raw["token_counts"]], + step_time_ms=[float(v) for v in raw["step_time_ms"]], + fixed_overhead_ms=float(raw.get("fixed_overhead_ms", 0.0)), + batch_sizes=[int(v) for v in raw.get("batch_sizes", [])], + batch_overhead_ms=[float(v) for v in raw.get("batch_overhead_ms", [])], + ) + else: + # Flat = "unprofiled": the planner keeps verifying the full block + # rather than trimming on a cost model where every token is free. + cost_table = SpsCostTable.flat() + logger.warning( + "DSpark confidence scheduling is enabled but no " + "confidence_sps_table_path was provided. Without a profiled step-cost " + "curve the planner cannot tell a cheap verify token from an expensive " + "one, so it will keep verifying the full block (no scheduling gain)." + ) + + # Calibration lives on the draft's confidence head; fall back to a plain + # sigmoid when the head is absent (no confidence to calibrate anyway). + from .dspark_sts import resolve_confidence_head + head = resolve_confidence_head(draft_model) + apply_calibration = head.apply_sts if head is not None else None + sts_path = self.spec_config.confidence_sts_path + if sts_path and head is None: + # Refuse rather than degrade: a configured calibration that cannot + # be attached is a wiring bug, never a preference. + raise ValueError( + f"confidence_sts_path is set but no confidence head was found " + f"on {type(draft_model).__name__}; calibration cannot be " + f"applied, so the planner would schedule on raw sigmoid " + f"survivals. This is a model-wiring bug, not a config choice.") + if sts_path and head is not None: + # Accepts either key spelling ("sts_temperatures" here, + # "temperatures" in SGLang); the vectors are interchangeable. + from .dspark_sts import load_sts_temperatures_from_path + temps = load_sts_temperatures_from_path(sts_path) + head.load_sts_temperatures(torch.tensor(temps, dtype=torch.float32)) + logger.info(f"DSpark: loaded STS calibration from {sts_path}") + + tiers = self.spec_config.verify_len_tiers + if not tiers: + # Derive the ladder from the cost curve. Each tier is a captured + # graph whose memory comes out of KV cache, so the derivation is + # capped; tiers are only zero-loss at the batch size they were + # derived for. A flat table yields [min, block_size] + # (no scheduling). + from .dspark_planner import derive_verify_len_tiers + + tiers = derive_verify_len_tiers( + cost_table=cost_table, + num_requests=max(int(max_batch), 1), + block_size=block_size, + min_verify_len=cfg.min_verify_len, + ) + logger.info(f"DSpark: derived verify-length tiers {tiers}") + self.verify_planner = DSparkVerifyPlanner( + cfg=cfg, + cost_table=cost_table, + tiers=tiers, + apply_calibration=apply_calibration, + ) + + # The env var overrides the config knob so the mode can be swapped + # without editing a serving config. + from .dspark_observability import (DSparkRaggedStats, RaggedVerifyMode, + read_ragged_verify_mode) + + configured = (RaggedVerifyMode.COMPACT + if self.spec_config.enable_ragged_verify else + RaggedVerifyMode.STATIC) + self.ragged_verify_mode = read_ragged_verify_mode(default=configured) + if (self.verify_planner is not None + and self.verify_planner.forced_budget_frac is not None + and not self.ragged_verify_mode.trims_submitted_tokens): + raise ValueError( + "TLLM_DSPARK_FORCE_BUDGET_FRAC is set but ragged verify mode " + f"'{self.ragged_verify_mode.value}' does not trim submitted " + "tokens; the fraction would be logged as FORCED and silently " + "never applied. Unset it or run compact mode.") + # STS collection is env-gated; it costs one device->host copy per + # decode step. + from .dspark_sts import make_recorder_from_env + self.sts_recorder = make_recorder_from_env( + block_size=int(block_size), + rank=int(self.mapping.rank), + has_cost_table=bool( + self.spec_config.confidence_sps_table_path), + ragged_mode=self.ragged_verify_mode.value) + + self.ragged_stats = DSparkRaggedStats(mode=self.ragged_verify_mode, + max_draft_len=block_size) + # Shared (not copied) so the summary stays live. + self.ragged_stats.planner_stats = self.verify_planner.stats + logger.info( + f"DSpark verify planner: tiers={self.verify_planner.tiers}, " + f"profiled_cost_table={not cost_table.is_flat}, " + f"ragged_verify_mode={self.ragged_verify_mode.value}" + ) + if (self.ragged_verify_mode.trims_submitted_tokens + and cost_table.is_flat): + logger.warning( + "DSpark ragged verify is set to 'compact' but the cost table is " + "flat, so the planner's budget degenerates to verify-all and " + "every request will receive the same full window. Profile a " + "cost table and pass confidence_sps_table_path, or the ragged " + "path will not be exercised." + ) + def _assign_slot(self, req_id: int, reset: bool) -> int: """Get (or refresh) the slot for a request; reset clears its window.""" if reset and req_id in self._req_to_slot: @@ -325,6 +568,11 @@ def _assign_slot(self, req_id: int, reset: bool) -> int: self._req_to_slot[req_id] = slot self._ctx_len[slot] = 0 self._kv_windows[slot].zero_() + if self._confidence_logits is not None: + # A recycled slot still holds the previous occupant's scores; + # reset to neutral so the first scheduling decision verifies + # the full block. + self._confidence_logits[slot].fill_(_NEUTRAL_CONFIDENCE_LOGIT) return self._req_to_slot[req_id] def _seed_context_windows( @@ -379,18 +627,22 @@ def _draft_gen_block_batched( ) -> torch.Tensor: """CUDA-graph-safe batched gen draft (all gen requests in one forward). - Free of host syncs and data-dependent shapes: per-request quantities - (``nacc``, the bonus, ``main_hidden``, ``start_pos``, the multi-accept - back-fill) are gathered as tensors, slots come from the host-built - ``_batch_to_slot`` mirror, and the backbone runs once via - ``DSparkDraftModel.forward_batched``. Returns the per-position corrected - block logits ``[num_gens, K, vocab]`` (or ``None`` when there is nothing to - draft); the worker feeds them to ``SpecWorkerBase.sample_draft_tokens``. - Confidence truncation stays disabled — the full block is proposed. + Must stay free of host syncs and data-dependent shapes. Returns block + logits ``[num_gens, K, vocab]`` for ``sample_draft_tokens``, or + ``None`` when there is nothing to draft. The full block is always + proposed. """ num_gens = batch_size - num_contexts + # K is the draft block width (fixed by the checkpoint); Kp1 is the + # target's per-gen-request token stride in ``captured`` + # (runtime_draft_len + 1, which shrinks when verification is trimmed). + # Striding by max_draft_len instead reads into the next request's + # hidden states. K = self.max_draft_len - Kp1 = K + 1 + # ``or K``: runtime_draft_len defaults to 0 (not None), so an + # ``is None`` check would leave a stride of 1 on unpopulated paths. + runtime_draft_len = spec_metadata.runtime_draft_len or K + Kp1 = int(runtime_draft_len) + 1 device = accepted_tokens.device if num_gens == 0: @@ -411,15 +663,25 @@ def _draft_gen_block_batched( accepted_tokens[num_contexts:batch_size].gather(1, gidx.unsqueeze(1)).squeeze(1).long() ) # [G] - # Captured target hidden at the bonus position within each request's Kp1 - # processed tokens. - arange_g = torch.arange(num_gens, device=device) - base = gen_start + arange_g * Kp1 # [G] - main_hidden = captured[base + gidx] # [G, ncap*hidden] - - # Fixed-size ([G, K]) masked back-fill of the intermediate accepted tokens - # (everything but the bonus) into the rolling window — same frames as the - # eager path (old+1 .. old+nacc-1), with j >= nacc-1 masked out. + # Start offset of each request's processed tokens inside ``captured``. + # Uniform batches stride by Kp1; under ragged verification the offsets + # must come from the same qo_indptr the input layout was packed with, + # or request r silently reads request r-1's tail. + qo_indptr = spec_metadata.qo_indptr + if spec_metadata.verify_lens is not None and qo_indptr is not None: + base = gen_start + qo_indptr[:num_gens].to(device=device, dtype=torch.long) + else: + arange_g = torch.arange(num_gens, device=device) + base = gen_start + arange_g * Kp1 # [G] + # Clamped like interim_base below: gidx comes from a device-side accept + # count, so a bad one must not turn into an out-of-bounds read. + main_hidden = captured[(base + gidx).clamp(min=0, max=captured.shape[0] - 1)] + + # Fixed-size [G, K] masked back-fill of the intermediate accepted + # tokens (all but the bonus) into frames old+1 .. old+nacc-1. The + # width stays at static K (capture-safe); ``nacc <= runtime_draft_len + # + 1`` bounds the valid j to this request's own block, so the extra + # columns are always masked out. old = self._ctx_len[slots] # [G] pre-increment decode position j = torch.arange(K, device=device) # [K] interim_valid = j.unsqueeze(0) < (nacc.unsqueeze(1) - 1) # [G, K] @@ -432,25 +694,30 @@ def _draft_gen_block_batched( interim_hidden, interim_pos, slots, interim_valid, self._kv_windows ) - # Advance the decode position by the accepted count; start_pos (= post- - # increment ctx_len) matches the eager path's frame value. + # Advance the decode position by the accepted count; start_pos is the + # post-increment ctx_len. start_pos = old + nacc # [G] self._ctx_len[slots] = start_pos - # Surface the per-position corrected block logits ([num_gens, K, vocab]) - # and let SpecWorkerBase.sample_draft_tokens do the (greedy or rejection) - # sampling + TP gather + draft_probs scatter, rather than argmaxing here. - _toks, _num_proposed, block_logits = draft_model.forward_batched( + _toks, confidence, block_logits = draft_model.forward_batched( main_hidden, bonus, start_pos, kv_windows=self._kv_windows, slots=slots, temperature=0.0, - confidence_threshold=0.0, + return_confidence=self.return_confidence, return_logits=True, all_rank_num_tokens=all_rank_num_tokens, ) + # Stash the [G, K] confidence for the verification scheduler. + # Slot-indexed (survives batch reshuffles), written in place (a + # captured graph keeps the same storage), and scattered through + # ``slots`` so a replay lands on the current batch's slots. + if confidence is not None: + self._confidence_logits[slots] = confidence.detach() + # Same slots, same capture region: content-only update, graph-safe. + self._confidence_stamp[slots] = self._draft_seq_cuda return block_logits def _forward_impl( @@ -476,36 +743,34 @@ def _forward_impl( spec_metadata._dspark_worker = self self._execute_guided_decoder_if_present(logits) - # Target-verify acceptance via the unified SpecWorkerBase entry: it - # reshapes the stored draft tokens (default (num_gens, runtime_draft_len) - # hook), then routes to strict or rejection sampling. Greedy parity with - # the previous hand-rolled path is preserved (rejection only engages for a - # non-greedy batch with valid draft_probs). + # Acceptance: strict target-verify, or rejection sampling for a + # non-greedy batch with valid draft_probs. accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( logits, attn_metadata, spec_metadata ) total_target_tokens = input_ids.shape[0] - # CUDA-graph warmup guard: the warmup forwards (is_cuda_graph set, stream - # NOT yet capturing) run synthetic gen batches that would otherwise advance - # the persistent rolling-window state. Snapshot and restore it so warmup is - # side-effect-free. (During the capture pass itself the stream IS capturing, - # so we skip the save/restore and let the ops be recorded; real requests - # reset their slot's window+ctx_len at prefill, wiping any capture-time - # mutation.) + # CUDA-graph warmup guard: warmup forwards (is_cuda_graph set, stream + # NOT yet capturing) run synthetic gen batches; snapshot/restore the + # persistent window state so they are side-effect-free. During the + # capture pass the stream IS capturing, so skip the save/restore and + # let the ops be recorded; prefill resets wipe capture-time mutation. is_warmup = ( - getattr(spec_metadata, "is_cuda_graph", False) + spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() ) if is_warmup: saved_ctx_len = self._ctx_len.clone() saved_windows = self._kv_windows.clone() + # The confidence buffer is persistent state too; warmup scores + # must not reach the verification planner. + saved_confidence = ( + None if self._confidence_logits is None else self._confidence_logits.clone() + ) - # Assign / reset window slots for context (prefill) requests and seed each - # request's rolling KV window from its prompt's captured context, so the - # first generation step drafts against real context instead of an all-zero - # window (acceptance-rate only; verified decoding keeps output correct). + # Seed prefill requests' rolling windows from their prompt's captured + # context (affects acceptance rate only). if num_contexts > 0: self._seed_context_windows( draft_model, @@ -515,24 +780,18 @@ def _forward_impl( total_target_tokens, ) - # FUSED_COMM MoE backends (DeepGEMM MegaMoE) synchronize EP ranks with an - # in-kernel phase-flip NVLink barrier that flips on every kernel call, so - # every rank must invoke the draft MoE the same number of times and with - # the same globally-gathered per-rank token list, or the barrier desyncs - # (hang / "unspecified launch failure"). The draft runs over generation - # requests only, each expanded to ``block`` positions, so the per-rank - # draft-MoE token count is ``num_gens * block``. ``all_rank_num_gens`` is - # gathered at metadata-prep time (model_engine, outside any CUDA-graph - # capture region); it is None for non-ADP / single-rank runs, where the - # local ``[num_tokens]`` fallback in ``_forward_stage`` is correct. + # FUSED_COMM MoE backends sync EP ranks with a phase-flip barrier: + # every rank must invoke the draft MoE the same number of times with + # the same globally-gathered per-rank token list (num_gens * block + # each) or the barrier desyncs. ``all_rank_num_gens`` is gathered at + # metadata-prep time, outside any capture region; None for non-ADP / + # single-rank runs (local fallback in ``_forward_stage``). block = int(draft_model.block_size) - all_rank_num_gens = getattr(spec_metadata, "all_rank_num_gens", None) - # A rank with zero local gen requests still has to cross the draft MoE's - # cross-rank barrier, but DeepseekV4MoE's router / shared-expert dense - # GEMMs reject a 0-row input (cuBLAS CUBLAS_STATUS_INVALID_VALUE), so such - # a rank runs a single 1-row dummy through the MoE (like ADP padding). - # Encode that as ``1`` in the globally-shared per-rank token list so every - # rank agrees on the FUSED_COMM chunk count and per-rank slice. + all_rank_num_gens = spec_metadata.all_rank_num_gens + # A rank with zero local gen requests must still cross the draft MoE's + # barrier, but the router / shared-expert GEMMs reject 0-row input, so + # it runs a 1-row dummy -- encoded as ``1`` in the shared per-rank list + # so every rank agrees on the chunk count. all_rank_draft_tokens = ( [max(1, int(g) * block) for g in all_rank_num_gens] if all_rank_num_gens is not None @@ -543,8 +802,6 @@ def _forward_impl( ) if num_gens > 0: - # The batched gen-block draft returns the per-position corrected block - # logits [num_gens, K, vocab] and is CUDA-graph-safe. gen_logits = self._draft_gen_block_batched( draft_model, spec_metadata, @@ -557,25 +814,20 @@ def _forward_impl( all_rank_num_tokens=all_rank_draft_tokens, ) if gen_logits is not None: - # SpecWorkerBase samples the draft tokens (greedy argmax, or - # rejection sampling for a non-greedy batch), performs the TP - # gather, and scatters the proposal distribution into draft_probs. gen_draft_tokens = self.sample_draft_tokens( gen_logits, spec_metadata, batch_size, num_contexts=num_contexts ) - # The context one-hot must match the width the gen scatter just - # published to draft_probs, NOT gen_logits.shape[-1]: under TP the - # draft logits are vocab-sharded and sample_draft_tokens gathers - # them to full vocab before scattering, so the pre-gather shard - # width would leave stale columns and corrupt rejection. + # The context one-hot must match the width published to + # draft_probs, NOT gen_logits.shape[-1]: under TP the draft + # logits are vocab-sharded before sample_draft_tokens gathers + # them to full vocab. gen_vocab = spec_metadata.draft_probs_last_dim else: gen_draft_tokens = torch.zeros((num_gens, K), dtype=torch.int32, device="cuda") gen_vocab = None else: - # No local generation requests: if any peer EP rank has some, we must - # still cross the draft MoE's cross-rank barrier the same number of - # times (zero-token) so a FUSED_COMM phase-flip barrier stays lockstep. + # No local gens: if any peer EP rank has some, still cross the + # draft MoE barrier so the FUSED_COMM phase-flip stays lockstep. if global_has_gen: draft_model.run_moe_lockstep_noop(all_rank_draft_tokens, accepted_tokens.device) gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") @@ -603,11 +855,33 @@ def _forward_impl( if is_warmup: self._ctx_len.copy_(saved_ctx_len) self._kv_windows.copy_(saved_windows) + if saved_confidence is not None: + self._confidence_logits.copy_(saved_confidence) - return { + outputs = { "logits": raw_logits, "new_tokens": accepted_tokens, "new_tokens_lens": num_accepted_tokens, "next_draft_tokens": next_draft_tokens, "next_new_tokens": next_new_tokens, } + # cap-accept only: per-request trim written by ``apply_accept_caps``. + # Rides the sampler's existing copy, so it costs no extra sync. + cap_trim = spec_metadata.accept_cap_trim + if cap_trim is not None and spec_metadata.accept_caps is not None: + outputs["cap_trim_lens"] = cap_trim[:batch_size] + # Ragged steps: the TOKEN windows the step actually verified, straight + # off the persistent layout buffer. Rides the sampler's existing D2H + # batch; under device-window selection this is the ONLY place the host + # can learn the true windows (py_verify_len then holds only the host + # shape split), and the rewind arithmetic must use it. + # + # Gen-only steps only: the buffer is generation-indexed while every + # other output row is batch-indexed, so on a mixed step (eager ragged + # with finished contexts) the rows would misalign AND come up short. + # Mixed steps keep exact windows on py_verify_len, so the sampler's + # snapshot fallback stays correct there. + if (spec_metadata.verify_lens is not None + and spec_metadata.verify_lens.shape[0] >= batch_size): + outputs["verify_lens"] = spec_metadata.verify_lens[:batch_size] + return outputs diff --git a/tensorrt_llm/_torch/speculative/dspark_device_select.py b/tensorrt_llm/_torch/speculative/dspark_device_select.py new file mode 100644 index 000000000000..a2febd3ff914 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_device_select.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Device-side verify-window selection: the in-graph prologue. + +The host planner decides windows from a confidence snapshot that is one +iteration old (the batch reshuffles across the lag, and the read itself must +not sync). This module ranks the block that is about to be verified by its +OWN confidence instead: everything from the slot-indexed gather to the +per-token row maps is a pure tensor function, so it can run at the head of +the captured graph, after the previous step's draft wrote the confidence +buffer and before this step's verify attention reads the layout. + +The split of responsibilities follows the paper's dual-timescale design: + +* the CAPACITY -- ``(padded_bs, bucket, budget)`` -- stays host-decided from + the lagged snapshot, because it is the CUDA-graph key and the attention-DP + agreement payload, both of which must exist before launch; +* the RANKING -- which requests win the budget -- is computed here from + fresh confidence, at zero staleness. + +Nothing in this module raises on data: a captured graph cannot branch. +Feasibility (the bucket bounds) is checkable from host-known scalars before +launch; per-row staleness degrades to the neutral row (verify everything for +that request), mirroring the host planner's fail-open semantics. +""" + +from dataclasses import dataclass +from typing import Callable, Optional, Tuple + +import torch + +from .dspark_ragged import (build_qo_indptr, build_row_maps_device, + fill_bucket_device) +from .dspark_schedule import (NEUTRAL_CONFIDENCE_LOGIT, DSparkScheduleConfig, + compute_survival, schedule_verify_lens_topk) + +__all__ = ["DeviceWindowResult", "select_windows_device"] + + +@dataclass +class DeviceWindowResult: + """Everything the ragged layout consumers need, all device-resident. + + Attributes: + verify_lens: ``[padded_bs]`` int32 token windows (bonus included), + summing to exactly ``graph_num_tokens``; pad rows carry their + fill. + qo_indptr: ``[padded_bs + 1]`` int32 exclusive prefix sum. + req_idx: ``[graph_num_tokens]`` int64 owning-row per packed token. + kv_correction: ``[graph_num_tokens]`` int32; composed with a + ``kv_lens`` gather it yields each token's KV extent + (``refresh_ragged_row_kv_lens``). + """ + + verify_lens: torch.Tensor + qo_indptr: torch.Tensor + req_idx: torch.Tensor + kv_correction: torch.Tensor + + +def select_windows_device( + *, + confidence_logits: torch.Tensor, + slot_idx: torch.Tensor, + num_real: torch.Tensor, + budget: torch.Tensor, + graph_num_tokens: int, + cfg: DSparkScheduleConfig, + apply_calibration: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + stamp: Optional[torch.Tensor] = None, + expected_stamp: Optional[torch.Tensor] = None, + pad_len: Optional[int] = None, +) -> DeviceWindowResult: + """Rank the batch by fresh confidence and pack it into the agreed bucket. + + Mirrors the host chain ``_gather_rows -> apply_calibration -> + compute_survival -> schedule_verify_lens_topk -> fill_bucket`` with the + lag removed: ``confidence_logits`` is the LIVE slot-indexed buffer the + previous draft pass scattered into, not a staged snapshot. + + Args: + confidence_logits: ``[num_slots, K]`` raw logits, slot-indexed. + slot_idx: ``[padded_bs]`` int64 buffer row per batch position; + entries at or beyond ``num_real`` are never read meaningfully + (point them at any valid row, e.g. 0). + num_real: 0-d integer tensor; rows past it are padding. + budget: 0-d integer tensor, verify tokens above the floor -- the + host-decided (lagged) capacity knob, written before launch. + graph_num_tokens: the captured token bucket (capture constant). + cfg: scheduling bounds; ``resolved_max_verify_len + 1`` is the + per-row token ceiling. + apply_calibration: logits -> probabilities; sigmoid when None. + stamp: optional ``[num_slots]`` last-writer stamps; with + ``expected_stamp`` (0-d), rows whose stamp mismatches fall back + to the neutral row (full survival: verify everything), the same + fail-open the host planner applies to unknown slots. + pad_len: when given, pad rows carry EXACTLY this many tokens and all + fill slack goes to real rows -- the host fit's published split, + which its pre-launch copy widths depend on (see + :func:`fill_bucket_device`). The caller must clamp ``budget`` so + the real rows can absorb ``graph_num_tokens - n_pad * pad_len``. + + Returns: + :class:`DeviceWindowResult`; every tensor is derived without a host + sync, so the call is capture-safe end to end. + """ + if confidence_logits.dim() != 2: + raise ValueError( + f"confidence_logits must be [num_slots, K], got " + f"{tuple(confidence_logits.shape)}") + device = confidence_logits.device + padded_bs = slot_idx.numel() + + selected = confidence_logits.index_select(0, slot_idx.to(torch.long)) + if stamp is not None: + if expected_stamp is None: + raise ValueError("stamp requires expected_stamp") + stale = stamp.index_select(0, slot_idx.to(torch.long)) != expected_stamp + selected = torch.where( + stale.unsqueeze(1), + torch.full_like(selected, NEUTRAL_CONFIDENCE_LOGIT), + selected, + ) + + calibrate = apply_calibration or torch.sigmoid + survival = compute_survival(calibrate(selected)) + # Pad rows must win no budget: zero survival sits below survival_eps, so + # the ranking never selects them and their windows stay at the floor + # until the fill tops them up. + is_real = torch.arange(padded_bs, device=device) < num_real + survival = torch.where(is_real.unsqueeze(1), survival, + torch.zeros_like(survival)) + + scheduled = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=cfg) + # The scheduler counts drafted positions; the token window adds the + # bonus/anchor, matching every host fill_bucket callsite. + token_lens = scheduled + 1 + max_token_len = int(cfg.resolved_max_verify_len) + 1 + filled = fill_bucket_device( + token_lens, + num_real=num_real, + graph_num_tokens=graph_num_tokens, + max_verify_len=max_token_len, + pad_fill=pad_len, + ) + req_idx, correction = build_row_maps_device( + filled, graph_num_tokens=graph_num_tokens) + return DeviceWindowResult( + verify_lens=filled, + qo_indptr=build_qo_indptr(filled), + req_idx=req_idx, + kv_correction=correction, + ) diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py new file mode 100644 index 000000000000..4fcce61c6016 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -0,0 +1,400 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Verification-mode selection and per-step counters for DSpark ragged verify. + +Failures on this path are silent (output stays correct at baseline accuracy), +so activity is counted rather than inferred. Modes: ``static`` verifies the full block; ``cap-accept`` computes and +commits per-request windows but submits the full block, so its output must +stay bit-identical to ``static``; ``compact`` submits only each request's +window and needs a profiled :class:`~.dspark_planner.SpsCostTable` -- with a +flat cost model the budget degenerates to verify-all. +""" + +import os +from collections import Counter +from enum import Enum +from typing import Dict, Optional, Sequence, Tuple + +from ...logger import logger + +__all__ = [ + "RaggedVerifyMode", + "read_ragged_verify_mode", + "DSparkRaggedStats", +] + +#: Overrides ``DSparkDecodingConfig``. Read per call, not cached at import: +#: tests flip it with ``monkeypatch.setenv``. +RAGGED_VERIFY_MODE_ENV = "TLLM_DSPARK_RAGGED_VERIFY_MODE" + + +class RaggedVerifyMode(str, Enum): + """How much of the drafted block reaches the target. See module docstring.""" + + STATIC = "static" + CAP_ACCEPT = "cap-accept" + COMPACT = "compact" + + @property + def computes_windows(self) -> bool: + """Whether per-request windows are computed at all (``cap-accept`` + shares every host-side path with ``compact`` except the token-axis + trim).""" + return self is not RaggedVerifyMode.STATIC + + @property + def trims_submitted_tokens(self) -> bool: + """Whether the token axis handed to the target actually shrinks.""" + return self is RaggedVerifyMode.COMPACT + + +def read_ragged_verify_mode( + default: RaggedVerifyMode = RaggedVerifyMode.STATIC +) -> RaggedVerifyMode: + """Resolve the verification mode from the environment. + + Raises: + ValueError: on an unrecognized value; a silent fallback to ``static`` + would turn a typo into "the feature never ran". + """ + value = os.environ.get(RAGGED_VERIFY_MODE_ENV) + if value is None or value == "": + return default + for mode in RaggedVerifyMode: + if value == mode.value: + return mode + raise ValueError( + f"invalid {RAGGED_VERIFY_MODE_ENV}={value!r}; expected one of " + f"{', '.join(repr(m.value) for m in RaggedVerifyMode)}") + + +def resolve_ragged_verify_mode(spec_config) -> RaggedVerifyMode: + """The mode a config plus the environment actually select. + + ``enable_ragged_verify`` alone is not the answer: ``cap-accept`` sets it + too but submits the FULL block, so batch-shaping code (capture set, graph + key) must ask :attr:`RaggedVerifyMode.trims_submitted_tokens`. Must stay + identical to the DSpark worker's own resolution. + """ + configured = (RaggedVerifyMode.COMPACT if getattr( + spec_config, "enable_ragged_verify", False) else RaggedVerifyMode.STATIC) + return read_ragged_verify_mode(default=configured) + + +def trims_submitted_tokens(spec_config) -> bool: + """Whether this config+environment shrinks the token axis. See above.""" + if spec_config is None: + return False + return resolve_ragged_verify_mode(spec_config).trims_submitted_tokens + + +class DSparkRaggedStats: + """Per-step counters for the ragged verification path. + + Host-only and cheap enough to leave on: no device reads, no + synchronization. + """ + + def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): + self.mode = mode + self.max_draft_len = int(max_draft_len) + #: Planner decision counters, attached by the worker once the planner + #: exists; the only record of WHY a step declined to trim + #: (`fallback_flat_cost` = no profiled SPS table, so no trimming). + self.planner_stats: Optional[Dict[str, int]] = None + + self.steps_total = 0 + #: Steps where per-request windows were produced AND differed; a step + #: where every request got the same window counts as uniform. + self.steps_ragged = 0 + #: Windows were produced, but every request got the same one. + self.steps_uniform_windows = 0 + #: No windows at all -- the planner declined or the shape did not fit. + self.steps_no_windows = 0 + + #: Tokens the target would have scored with no trimming at all. + self.ceiling_tokens = 0 + #: Tokens the windows asked for, before bucket rounding. + self.window_tokens = 0 + #: Tokens actually submitted, after rounding up to a captured bucket. + self.delivered_tokens = 0 + + self.verify_len_hist: Counter = Counter() + self.bucket_hist: Counter = Counter() + self.padded_bs_hist: Counter = Counter() + #: why the graph runner declined -> count. "peer_shape_mismatch" = one + #: rank declined ragged and dragged the world eager; "key_not_captured" + #: = wrong bucket grid; "peer_not_gen_only" = ordinary continuous + #: batching. + #: Steps whose batch held at least two generation requests; the gate's + #: floor counts these, not all steps (concurrency ramps up). + self.steps_multi_request = 0 + self.graph_miss_reasons: Dict[str, int] = {} + #: the actual graph key, kept only for the reasons where it localizes + #: the bug. + self.graph_miss_shapes: Dict[str, int] = {} + + #: Acceptance, split by whether the request's window was shortened. + self.accepted_tokens = 0 + self.requests_scored = 0 + #: Requests whose window was shorter than the full block. + self.requests_trimmed = 0 + #: Trimmed requests that accepted their ENTIRE window: the draft was + #: still alive at the cut, so the trim certainly cost acceptance (one + #: that accepted fewer would have died on its own and cost nothing). + self.trimmed_hit_ceiling = 0 + #: Positions the target accepted and the window then discarded. Only + #: ``cap-accept`` scores these; under ``compact`` 0 means "unknowable", + #: not "no loss". + self.cap_trim_tokens = 0 + #: Kept per request because a broad mild trim and a concentrated one + #: sum to the same total and mean opposite things about the planner. + self.requests_cap_trimmed = 0 + self.cap_trim_max = 0 + self.cap_trim_hist: Counter = Counter() + + #: An eager step costs far more than the tokens the trim saved. + self.graph_replays = 0 + self.graph_eager = 0 + + #: Reasons the step declined to go ragged, keyed by a short slug. + self.fallbacks: Counter = Counter() + + def record_step( + self, + *, + num_gen_requests: int, + verify_lens: Optional[Sequence[int]] = None, + bucket: Optional[int] = None, + padded_bs: Optional[int] = None, + fallback: Optional[str] = None, + delivered: Optional[int] = None, + ) -> None: + """Record one decode step's scheduling decision. + + Args: + num_gen_requests: generation requests in the batch. + verify_lens: per-request drafted positions verified, or None when + the step stayed uniform. Excludes the bonus token, matching + ``py_verify_len``. + bucket: the captured token bucket the batch was fitted to, if any. + padded_bs: the row count the batch was padded to, if any. + fallback: short reason the step did not go ragged. + delivered: tokens actually handed to the target, when that is not + derivable from ``verify_lens`` and ``bucket``. Required by + ``cap-accept``, which schedules windows but submits the full + block. + """ + self.steps_total += 1 + # A batch of one is uniform by construction; counting it toward the + # gate's floor would make the gate fire during ramp-up. + if int(num_gen_requests) >= 2: + self.steps_multi_request += 1 + # Ceiling = what a no-trim step would have submitted: each ROW sends + # its bonus token plus the full block. ROWS, not requests: on a fitted + # bucket `delivered` books `padded_bs * (tier + 1)` where `padded_bs` + # is the CUDA-graph-rounded max row count ACROSS RANKS, so the ceiling + # must count the same rows or trim_ratio goes negative. Only that case + # rebases; every other path books `delivered` over the real rows, and + # widening their ceiling would understate the trim. + ceiling_rows = int(num_gen_requests) + if bucket and delivered is None and padded_bs: + ceiling_rows = max(ceiling_rows, int(padded_bs)) + self.ceiling_tokens += ceiling_rows * (1 + self.max_draft_len) + + if fallback: + self.fallbacks[fallback] += 1 + + if not verify_lens: + self.steps_no_windows += 1 + self.window_tokens += int(num_gen_requests) * (1 + + self.max_draft_len) + self.delivered_tokens += int(num_gen_requests) * ( + 1 + self.max_draft_len) + return + + lens = [int(v) for v in verify_lens] + self.verify_len_hist.update(lens) + self.window_tokens += sum(1 + v for v in lens) + # A bucket is what the graph actually ran; without one the step ran on + # exactly the windows it asked for -- unless the caller says otherwise, + # which is how cap-accept reports submitting the full block. + if delivered is not None: + self.delivered_tokens += int(delivered) + else: + self.delivered_tokens += int(bucket) if bucket else sum(1 + v + for v in + lens) + if bucket: + self.bucket_hist[int(bucket)] += 1 + if padded_bs: + self.padded_bs_hist[int(padded_bs)] += 1 + + if len(set(lens)) > 1: + self.steps_ragged += 1 + else: + self.steps_uniform_windows += 1 + + def record_graph(self, *, replayed: bool, shape) -> None: + """Record whether this step replayed a CUDA graph or ran eager. + + ``shape`` is the ``(num_rows, total_verify_tokens)`` the step actually + submitted, or None on a replay; kept only for misses. Deliberately + required, not defaulted: a default would make "probe not wired" + indistinguishable from "no misses had shapes". + """ + if replayed: + self.graph_replays += 1 + return + self.graph_eager += 1 + if shape is not None: + reason, key = shape + self.graph_miss_reasons[reason] = ( + self.graph_miss_reasons.get(reason, 0) + 1) + if key is not None: + text = str(key) + self.graph_miss_shapes[text] = ( + self.graph_miss_shapes.get(text, 0) + 1) + + def merge_planner_stats(self, planner_stats: Dict[str, int]) -> None: + """Fold in the planner's own fallback counters.""" + for key, value in planner_stats.items(): + if key.startswith("fallback_") and value: + self.fallbacks[key] = int(value) + + @property + def trim_ratio(self) -> float: + """Fraction of ceiling tokens the scheduler actually removed. + + ``0.0`` means the feature ran but saved nothing, which is the signature + of an unprofiled cost table. + """ + if self.ceiling_tokens <= 0: + return 0.0 + return 1.0 - (self.delivered_tokens / self.ceiling_tokens) + + @property + def distinct_verify_lens(self) -> int: + """How many different window sizes were ever handed out.""" + return len(self.verify_len_hist) + + def record_acceptance(self, + *, + accepted: int, + window: int, + cap_trim: int = 0) -> None: + """Record one request's acceptance against the window it was given. + + Args: + accepted: drafted positions accepted (excludes the bonus token). + window: drafted positions the request was allowed to verify. + cap_trim: positions this request accepted that its window then + discarded. Non-zero only under ``cap-accept``, the one mode + that scores them; elsewhere 0 means "not measured", not "no + loss". + """ + self.requests_scored += 1 + self.accepted_tokens += int(accepted) + if int(window) < self.max_draft_len: + self.requests_trimmed += 1 + if int(accepted) >= int(window): + self.trimmed_hit_ceiling += 1 + cap_trim = int(cap_trim) + if cap_trim > 0: + self.cap_trim_tokens += cap_trim + self.requests_cap_trimmed += 1 + self.cap_trim_hist[cap_trim] += 1 + if cap_trim > self.cap_trim_max: + self.cap_trim_max = cap_trim + + @property + def accept_len(self) -> float: + """Mean accepted drafted positions per request. 0.0 with no data.""" + if not self.requests_scored: + return 0.0 + return self.accepted_tokens / self.requests_scored + + @property + def accept_loss_per_request(self) -> float: + """Mean accepted positions the schedule discarded, per request. + + :attr:`accept_len` plus this is the no-trim acceptance on the same + drafts. Meaningful only under ``cap-accept``; elsewhere 0.0 means "not + measured", not "no loss". + """ + if not self.requests_scored: + return 0.0 + return self.cap_trim_tokens / self.requests_scored + + @property + def cap_trim_concentration(self) -> float: + """Share of scored requests that lost anything at all. ``cap-accept``. + + Distinguishes a broad mild trim from the same mean loss concentrated + on a few requests; only the second is a problem. + """ + if not self.requests_scored: + return 0.0 + return self.requests_cap_trimmed / self.requests_scored + + @property + def trim_regret_rate(self) -> float: + """Share of trimmed requests whose draft was still alive at the cut. + + 0.0 means every trim was free (those drafts would have died anyway); + toward 1.0 the trims are buying throughput with acceptance. + """ + if not self.requests_trimmed: + return 0.0 + return self.trimmed_hit_ceiling / self.requests_trimmed + + def summary(self) -> Dict[str, object]: + """A JSON-safe snapshot, for logging or assertions.""" + return { + "mode": self.mode.value, + "steps_total": self.steps_total, + "steps_ragged": self.steps_ragged, + "steps_multi_request": self.steps_multi_request, + "steps_uniform_windows": self.steps_uniform_windows, + "steps_no_windows": self.steps_no_windows, + "ceiling_tokens": self.ceiling_tokens, + "window_tokens": self.window_tokens, + "delivered_tokens": self.delivered_tokens, + "trim_ratio": round(self.trim_ratio, 4), + "distinct_verify_lens": self.distinct_verify_lens, + "verify_len_hist": dict(sorted(self.verify_len_hist.items())), + "bucket_hist": dict(sorted(self.bucket_hist.items())), + "padded_bs_hist": dict(sorted(self.padded_bs_hist.items())), + "graph_miss_reasons": dict(sorted(self.graph_miss_reasons.items())), + "graph_miss_keys": dict(sorted(self.graph_miss_shapes.items())), + "accept_len": round(self.accept_len, 4), + "requests_scored": self.requests_scored, + "requests_trimmed": self.requests_trimmed, + "trim_regret_rate": round(self.trim_regret_rate, 4), + "cap_trim_tokens": self.cap_trim_tokens, + "accept_loss_per_request": round(self.accept_loss_per_request, 4), + "requests_cap_trimmed": self.requests_cap_trimmed, + "cap_trim_max": self.cap_trim_max, + "cap_trim_hist": dict(sorted(self.cap_trim_hist.items())), + "cap_trim_concentration": round(self.cap_trim_concentration, 4), + "graph_replays": self.graph_replays, + "graph_eager": self.graph_eager, + "fallbacks": dict(self.fallbacks), + "planner": dict(self.planner_stats) if self.planner_stats else {}, + } + + def log_summary(self, *, prefix: str = "DSpark ragged verify") -> None: + summary = self.summary() + logger.info(f"{prefix}: {summary}") diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py new file mode 100644 index 000000000000..17b4a100108c --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -0,0 +1,360 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DSpark verification-budget planner: how many tokens are worth verifying. + +The scheduler maximizes *system* goodput:: + + tau(n) = num_gen_requests + sum of the n largest prefix-survival values + Theta(n) = tau(n) / step_cost_ms(num_gen_requests + n) + budget = argmax_n Theta(n) + +Invariants: take the GLOBAL argmax (marginal cost is a staircase, so ``Theta`` +is not unimodal and a first-descent stop parks on the first shelf); cost is +per step, indexed by the TOTAL verified token count including the bonus +position (a request verifying ``L`` drafts submits ``L + 1`` tokens); and +``Theta`` is a ratio, so the non-trimmable ``bias + alpha(bs)`` terms do not +cancel -- they must be supplied, not guessed (see :class:`SpsCostTable`). The +planner runs on the host because its output selects a CUDA graph, a host +decision. +""" + +import math +from dataclasses import dataclass, field +from typing import List, Optional, Sequence + +import numpy as np + +from ...logger import logger + +__all__ = [ + "SpsCostTable", + "check_table_fingerprint", + "compute_verify_token_budget", + "budget_argmax_over_uniform_lens", + "derive_verify_len_tiers", + "total_verify_tokens", +] + + +def total_verify_tokens(num_requests: int, verify_len: int) -> int: + """Tokens the target scores for ``num_requests`` each verifying ``verify_len`` drafts. + + One bonus/anchor token per request plus its drafted positions. This is the + unit :class:`SpsCostTable` is indexed by; the single place it is computed so + the planner and the tier derivation cannot drift apart. + """ + return int(num_requests) * (int(verify_len) + 1) + + +@dataclass(frozen=True) +class SpsCostTable: + """Measured decode step time as a function of total verified tokens. + + Attributes: + token_counts: strictly increasing total-token breakpoints, in + :func:`total_verify_tokens` units (bonus tokens included). + step_time_ms: measured ``theta(M)`` at each breakpoint. + fixed_overhead_ms: the ``bias`` term -- per-step cost that trimming + cannot touch -- if not already inside ``step_time_ms``. + Understating the non-trimmable terms makes the planner over-trim: + ``tau / T`` is a ratio, so they do not cancel. + batch_sizes / batch_overhead_ms: optional ``alpha(bs)`` breakpoints -- + the batch-size-dependent, non-trimmable part of the step. + + Lookup is a clamped linear interpolation between breakpoints; a flat shelf + must be MEASURED as two breakpoints with equal values (a floor lookup makes + tokens look free right before a riser and the argmax over-spends). Queries + outside the measured range clamp to the end values; callers must bound the + budget so the high clamp does not happen silently. + """ + + token_counts: Sequence[int] + step_time_ms: Sequence[float] + fixed_overhead_ms: float = 0.0 + batch_sizes: Sequence[int] = field(default_factory=tuple) + batch_overhead_ms: Sequence[float] = field(default_factory=tuple) + + def __post_init__(self) -> None: + if len(self.token_counts) != len(self.step_time_ms): + raise ValueError( + f"token_counts ({len(self.token_counts)}) and step_time_ms " + f"({len(self.step_time_ms)}) must have the same length" + ) + if not self.token_counts: + raise ValueError("SpsCostTable requires at least one measured point") + if any(b <= a for a, b in zip(self.token_counts, self.token_counts[1:])): + raise ValueError("token_counts must be strictly increasing") + if any(not math.isfinite(t) or t <= 0.0 for t in self.step_time_ms): + raise ValueError("step_time_ms entries must be positive and finite") + if not math.isfinite(self.fixed_overhead_ms) or self.fixed_overhead_ms < 0.0: + raise ValueError("fixed_overhead_ms must be >= 0 and finite") + if len(self.batch_sizes) != len(self.batch_overhead_ms): + raise ValueError( + f"batch_sizes ({len(self.batch_sizes)}) and batch_overhead_ms " + f"({len(self.batch_overhead_ms)}) must have the same length" + ) + if any(b <= a for a, b in zip(self.batch_sizes, self.batch_sizes[1:])): + raise ValueError("batch_sizes must be strictly increasing") + if any(not math.isfinite(t) or t < 0.0 for t in self.batch_overhead_ms): + raise ValueError("batch_overhead_ms entries must be >= 0 and finite") + + def batch_overhead(self, num_requests: int) -> float: + """``alpha(num_requests)`` -- 0.0 when the table has no batch axis.""" + if not self.batch_sizes: + return 0.0 + return float( + np.interp( + float(num_requests), + np.asarray(self.batch_sizes, dtype=np.float64), + np.asarray(self.batch_overhead_ms, dtype=np.float64), + )) + + def step_time(self, num_tokens: int, num_requests: int = 0) -> float: + """``bias + alpha(num_requests) + theta(num_tokens)``, in ms.""" + return float(self.step_times(np.asarray([int(num_tokens)]), num_requests)[0]) + + def step_times(self, num_tokens: np.ndarray, num_requests: int = 0) -> np.ndarray: + """Vectorized :meth:`step_time` over a token-count array.""" + theta = np.interp( + np.asarray(num_tokens, dtype=np.float64), + np.asarray(self.token_counts, dtype=np.float64), + np.asarray(self.step_time_ms, dtype=np.float64), + ) + return theta + self.fixed_overhead_ms + self.batch_overhead(num_requests) + + @classmethod + def flat(cls, step_time_ms: float = 1.0) -> "SpsCostTable": + """Explicit "not profiled" marker: every extra token looks free, so + the planner degenerates to verify-all. Callers should warn rather than + silently ship this.""" + return cls(token_counts=(0,), step_time_ms=(float(step_time_ms),)) + + @property + def is_flat(self) -> bool: + """Whether the trimmable term ``theta(M)`` is constant. ``bias`` and + ``alpha(bs)`` are identical for every candidate length, so variation + there cannot distinguish lengths.""" + return len(set(self.step_time_ms)) <= 1 + + +def check_table_fingerprint(*, payload: dict, live: dict) -> None: + """Refuse a cost table profiled on a different engine than this one. + + A mismatched table produces no error downstream -- the planner just trims + to the wrong economics. Only the keys PRESENT in both dicts are compared, + so a table may carry facts the consumer cannot see (for humans to check), + and a table with no fingerprint loads with a warning rather than breaking + existing deployments. + """ + fp = (payload.get("_meta") or {}).get("engine") or payload.get("engine") + if not fp: + logger.warning( + "DSpark cost table carries no engine fingerprint; cannot verify it " + "was profiled on this configuration. Tables are engine-specific " + "(same cell measured 23%% apart across two configs) -- regenerate " + "with a current profiler to get load-time checking.") + return + def _facts_differ(a, b) -> bool: + try: + return float(a) != float(b) + except (TypeError, ValueError): + return str(a).upper() != str(b).upper() + + mismatched = { + key: (fp[key], live[key]) + for key in sorted(set(fp) & set(live)) + if _facts_differ(fp[key], live[key]) + } + if mismatched: + detail = ", ".join(f"{k}: table={a!r} engine={b!r}" + for k, (a, b) in mismatched.items()) + raise ValueError( + f"DSpark cost table was profiled on a different engine " + f"configuration ({detail}). Its step times do not describe this " + f"engine, so the planner would trim to the wrong economics. " + f"Re-profile on this configuration, or remove " + f"confidence_sps_table_path to run without trimming.") + unchecked = sorted(set(fp) - set(live)) + if unchecked: + logger.info( + f"DSpark cost table fingerprint: verified " + f"{sorted(set(fp) & set(live))}; not verifiable at this site " + f"(check manually): " + + ", ".join(f"{k}={fp[k]!r}" for k in unchecked)) + + +def compute_verify_token_budget( + *, + survival: np.ndarray, + num_gen_requests: int, + cost_table: SpsCostTable, + min_verify_len: int = 1, + max_verify_len: Optional[int] = None, + allowed_lens: Optional[Sequence[int]] = None, +) -> int: + """Return the batch-wide verify-token budget (tokens *above* the floor). + + Args: + survival: ``[bs, K]`` prefix-survival probabilities. Only the columns at + or beyond ``min_verify_len`` are schedulable; earlier positions are + already granted to every request. + num_gen_requests: generating requests this step. Each contributes one + guaranteed output token, which is the constant term of ``tau``. + cost_table: measured step cost; see :class:`SpsCostTable`. + min_verify_len / max_verify_len: per-request bounds, mirroring + :class:`~.dspark_schedule.DSparkScheduleConfig`. + allowed_lens: restrict the answer to budgets a captured tier can + realise (``n in {bs*(t - min_verify_len)}``). Pass it whenever the + result will be executed: any non-tier total is rounded UP to a + captured bucket, pushing a budget chosen inside a cost shelf back + over the riser it was avoiding. + + Returns: + ``n`` maximizing ``tau(n) / step_cost(floor_tokens + n)``, in ``[0, N]`` + where ``N`` is the number of schedulable candidates -- restricted to the + tier-aligned indices when ``allowed_lens`` is given. + """ + if survival.ndim != 2: + raise ValueError(f"survival must be [bs, K], got shape {survival.shape}") + bs, block_size = survival.shape + cap = min(int(max_verify_len or block_size), block_size) + schedulable = max(cap - int(min_verify_len), 0) + if bs == 0 or num_gen_requests <= 0 or schedulable <= 0: + return 0 + + surv = np.asarray(survival, dtype=np.float64) + candidates = np.sort(surv[:, min_verify_len : min_verify_len + schedulable].reshape(-1))[::-1] + if candidates.size == 0: + return 0 + + # tau(n): one bonus token per request, plus the floor positions' survival + # (unbought but part of the yield -- omitting the constant moves the argmax + # of a ratio), plus the n admitted candidates best-first. + base = float(num_gen_requests) + float(surv[:, :min_verify_len].sum()) + tau = base + np.concatenate(([0.0], np.cumsum(candidates))) + # Includes the bonus position: dropping it under-reports the batch by a + # whole ``bs``. + floor_tokens = total_verify_tokens(bs, min_verify_len) + tokens = floor_tokens + np.arange(tau.size) + theta = tau / cost_table.step_times(tokens, num_gen_requests) + if allowed_lens is None: + return int(np.argmax(theta)) + # Tier-aligned: the same theta curve, evaluated only at totals the executor + # can land on -- n = bs*(t - min_verify_len) gives tokens = bs*(t+1) + # = total_verify_tokens(bs, t). + idx = sorted({ + int(bs) * (int(t) - int(min_verify_len)) + for t in allowed_lens + if int(min_verify_len) <= int(t) <= cap + }) + idx = [n for n in idx if 0 <= n < theta.size] + if not idx: + return 0 + return int(max(idx, key=lambda n: theta[n])) + + +def derive_verify_len_tiers( + *, + cost_table: SpsCostTable, + num_requests: int, + block_size: int, + min_verify_len: int = 1, + max_tiers: int = 3, +) -> List[int]: + """Derive the verify lengths worth capturing a CUDA graph for. + + Within a measured flat cost shelf (adjacent breakpoints with equal values) + ``Theta = tau / cost`` strictly increases, so at a fixed batch size the + optimum sits on a shelf's *right edge*, never in its interior; this + returns those edges, always including ``min_verify_len`` and the full + block. On segments that genuinely rise the right-edge set is a + slope-change heuristic, and the property does not survive across batch + sizes (shelves live in token space) -- derive at the modal batch size and + measure the residual loss at the others. + + Args: + num_requests: batch size the tiers are derived for; tiers are a + function of it, so a deployment that captures several batch sizes + derives a set per batch size. + max_tiers: cap on how many lengths to return. Each extra tier is + another captured graph, which consumes memory that would otherwise + be KV cache. The full block and ``min_verify_len`` are kept first. + """ + lo, hi = int(min_verify_len), int(block_size) + if num_requests <= 0 or hi < lo: + return [lo] + + edges = set() + for breakpoint_tokens in cost_table.token_counts[1:]: + # Largest length whose total (num_requests * (L + 1), see + # total_verify_tokens) still sits below this breakpoint; inverting + # num_requests * L instead puts the edge one shelf too far right. + length = (int(breakpoint_tokens) - 1) // int(num_requests) - 1 + if lo <= length <= hi: + edges.add(length) + # The last shelf extends past the final breakpoint, so the whole block is + # always a right edge; the floor is always runnable. + edges.update({lo, hi}) + + tiers = sorted(edges) + if len(tiers) <= max_tiers: + return tiers + # Keep the endpoints; spread the remaining slots over the interior edges. + interior = tiers[1:-1] + step = len(interior) / float(max_tiers - 2) if max_tiers > 2 else 0 + keep = [interior[int(i * step)] for i in range(max_tiers - 2)] if max_tiers > 2 else [] + return sorted({tiers[0], *keep, tiers[-1]}) + + +def budget_argmax_over_uniform_lens( + *, + survival: np.ndarray, + num_gen_requests: int, + cost_table: SpsCostTable, + allowed_lens: Sequence[int], + min_verify_len: int = 1, +) -> int: + """Pick the best verify length from a discrete set (e.g. captured graph tiers). + + Evaluates ``Theta`` directly at each runnable length: optimizing the + continuous budget and rounding is not equivalent, since the rounded point + can land on the far side of a cost riser. + + Args: + allowed_lens: verify lengths with a captured graph. Values below + ``min_verify_len`` are ignored. + Returns: + The chosen per-request verify length. Falls back to ``min_verify_len`` + when no candidate is admissible. + """ + if survival.ndim != 2: + raise ValueError(f"survival must be [bs, K], got shape {survival.shape}") + bs, block_size = survival.shape + lens: List[int] = sorted( + {int(v) for v in allowed_lens if min_verify_len <= int(v) <= block_size} + ) + if bs == 0 or num_gen_requests <= 0 or not lens: + return int(min_verify_len) + + surv = np.asarray(survival, dtype=np.float64) + best_len, best_theta = lens[0], -np.inf + for length in lens: + # Expected yield: the bonus token plus survival of every verified + # position, floor included. + tau = float(num_gen_requests) + float(surv[:, :length].sum()) + theta = tau / cost_table.step_time(total_verify_tokens(bs, length), num_gen_requests) + if theta > best_theta: + best_len, best_theta = length, theta + return int(best_len) diff --git a/tensorrt_llm/_torch/speculative/dspark_ragged.py b/tensorrt_llm/_torch/speculative/dspark_ragged.py new file mode 100644 index 000000000000..8031b5d0e1ef --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_ragged.py @@ -0,0 +1,557 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Ragged verification layout: per-request verify lengths in one packed batch, +as a flat token axis plus a ``qo_indptr`` exclusive prefix sum. + +Two invariants keep the pipeline safe: ``verify_lens >= 1`` for every request +(position 0 carries the bonus/anchor token; a request that verifies nothing +never progresses), and ``sum(verify_lens) == graph_num_tokens`` after +:meth:`RaggedVerifyLayout.fill_bucket` -- ``seq_lens.sum()`` becomes +``attn_metadata.num_tokens``, which is all-gathered across attention-DP ranks +and drives the MoE's chunk count, and an under-filled bucket desynchronizes +them without raising. +""" + +import bisect +from dataclasses import dataclass +from typing import List, Optional, Sequence, Tuple + +import torch + +__all__ = [ + "RaggedVerifyLayout", + "build_qo_indptr", + "round_up_to_bucket", + "row_ids_from_lens", + "scatter_ragged_to_padded", + "count_accepted_ragged", + "fill_padded_rows_onehot", + "ragged_gather_index_lists", + "RaggedCaptureShape", + "choose_ragged_capture_shape", + "fill_bucket_device", + "build_row_maps_device", +] + + +def round_up_to_bucket(total: int, buckets: Sequence[int]) -> int: + """Smallest captured bucket that fits ``total``. + + Raises rather than clamping; callers must reject an oversized batch before + selecting a layout. + """ + if not buckets: + raise ValueError("round_up_to_bucket requires a non-empty bucket list") + if total > buckets[-1]: + raise ValueError( + f"total {total} exceeds the largest captured bucket {buckets[-1]}; " + f"the caller must reject this batch before building a layout" + ) + return buckets[bisect.bisect_left(buckets, total)] + + +def build_qo_indptr(verify_lens: torch.Tensor) -> torch.Tensor: + """``[bs]`` lengths -> ``[bs + 1]`` exclusive prefix sum, on the same device. + + Kept as a tensor op with no host sync so it can run inside a captured graph. + """ + if verify_lens.dim() != 1: + raise ValueError(f"verify_lens must be 1-D, got {tuple(verify_lens.shape)}") + lens = verify_lens.to(torch.int32) + out = torch.zeros(lens.numel() + 1, dtype=torch.int32, device=lens.device) + torch.cumsum(lens, dim=0, out=out[1:]) + return out + + +@dataclass +class RaggedVerifyLayout: + """One packed ragged verify batch. + + Attributes: + verify_lens: ``[bs]`` int32, each >= ``min_verify_len``. + qo_indptr: ``[bs + 1]`` int32 exclusive prefix sum of ``verify_lens``. + extend_start_loc: ``[bs]`` int32, ``qo_indptr[:-1]``. + graph_num_tokens: the captured token bucket this batch is padded to. + total_verify_tokens: ``sum(verify_lens)``, host-side when known. + num_real_requests: how many leading rows are real requests. Rows beyond + it are padding added to reach the captured batch size; their outputs + must be discarded. Equals ``bs`` until :meth:`fill_bucket` pads. + """ + + verify_lens: torch.Tensor + qo_indptr: torch.Tensor + extend_start_loc: torch.Tensor + graph_num_tokens: int + total_verify_tokens: Optional[int] = None + num_real_requests: Optional[int] = None + + @property + def bs(self) -> int: + return int(self.verify_lens.shape[0]) + + @property + def num_pad_requests(self) -> int: + return 0 if self.num_real_requests is None else self.bs - self.num_real_requests + + @classmethod + def from_verify_lens( + cls, + verify_lens: torch.Tensor, + *, + graph_num_tokens: Optional[int] = None, + buckets: Optional[Sequence[int]] = None, + total_verify_tokens: Optional[int] = None, + ) -> "RaggedVerifyLayout": + """Pack ``[bs]`` per-request lengths, rounding up to a captured bucket. + + Give ``graph_num_tokens`` directly, or ``buckets`` plus a host-side + ``total_verify_tokens``; deriving the total from the device tensor + would sync. + """ + lens = verify_lens.to(torch.int32) + if graph_num_tokens is None: + if buckets is None or total_verify_tokens is None: + raise ValueError( + "give graph_num_tokens, or buckets + total_verify_tokens; " + "deriving the total from the device tensor would sync" + ) + graph_num_tokens = round_up_to_bucket(int(total_verify_tokens), buckets) + indptr = build_qo_indptr(lens) + return cls( + verify_lens=lens, + qo_indptr=indptr, + extend_start_loc=indptr[:-1], + graph_num_tokens=int(graph_num_tokens), + total_verify_tokens=( + int(total_verify_tokens) if total_verify_tokens is not None else None + ), + ) + + def fill_bucket( + self, *, max_verify_len: int, padded_bs: Optional[int] = None + ) -> "RaggedVerifyLayout": + """Pad rows to ``padded_bs`` and tokens to exactly ``graph_num_tokens`` + (a short batch desynchronizes attention from the MoE; see the module + docstring). Slack goes to real requests first (up to ``max_verify_len``), + then to pad rows; every pad row needs at least one token or its empty + ``qo_indptr`` range breaks per-row slicing. Raises when the bucket + cannot be hit exactly. + """ + if self.total_verify_tokens is None: + raise ValueError("fill_bucket needs the host-side total_verify_tokens") + real: List[int] = [int(v) for v in self.verify_lens.tolist()] + n_real = len(real) + padded_bs = n_real if padded_bs is None else int(padded_bs) + if padded_bs < n_real: + raise ValueError( + f"padded_bs {padded_bs} < real request count {n_real}: the " + f"captured batch size cannot hold this batch" + ) + n_pad = padded_bs - n_real + + # Every pad row must carry at least one token. + baseline = sum(real) + n_pad + capacity = padded_bs * max_verify_len + if baseline > self.graph_num_tokens: + raise ValueError( + f"bucket {self.graph_num_tokens} is too small: {sum(real)} real " + f"tokens plus {n_pad} pad rows already needs {baseline}" + ) + if self.graph_num_tokens > capacity: + raise ValueError( + f"bucket {self.graph_num_tokens} exceeds what {padded_bs} rows can " + f"absorb at max_verify_len={max_verify_len} ({capacity}); pick a " + f"denser bucket grid or a larger padded_bs" + ) + + lens = real + [1] * n_pad + spare = self.graph_num_tokens - baseline + # Round-robin: spread extra tokens across real rows (not one request's + # low-survival tail), then spill to pad rows. + for lo, hi in ((0, n_real), (n_real, padded_bs)): + i = lo + while spare > 0 and any(lens[j] < max_verify_len for j in range(lo, hi)): + if lens[i] < max_verify_len: + lens[i] += 1 + spare -= 1 + i = lo + (i - lo + 1) % max(hi - lo, 1) + assert spare == 0, f"internal: {spare} tokens left after filling {lens}" + + filled = torch.tensor(lens, dtype=torch.int32, device=self.verify_lens.device) + out = RaggedVerifyLayout.from_verify_lens( + filled, + graph_num_tokens=self.graph_num_tokens, + total_verify_tokens=sum(lens), + ) + out.num_real_requests = n_real + return out + + def validate( + self, + *, + min_verify_len: int = 1, + max_verify_len: Optional[int] = None, + exact_fill: bool = False, + ) -> None: + """Assert the layout invariants (host-syncing: tests/debug only, never + the hot path). + + ``exact_fill`` additionally requires the token count to hit the bucket + exactly -- the post-:meth:`fill_bucket` contract for anything heading + into a captured graph; a short batch raises nowhere downstream. + """ + lens = [int(v) for v in self.verify_lens.tolist()] + if not lens: + raise ValueError("a ragged layout needs at least one request") + if min(lens) < max(min_verify_len, 1): + raise ValueError( + f"every request must verify its anchor token " + f"(verify_len >= {max(min_verify_len, 1)}), got {lens}" + ) + if max_verify_len is not None and max(lens) > max_verify_len: + raise ValueError( + f"verify_len exceeds the drafted block: {max(lens)} > {max_verify_len}" + ) + total = sum(lens) + if self.total_verify_tokens is not None and total != self.total_verify_tokens: + raise ValueError(f"total_verify_tokens {self.total_verify_tokens} != sum {total}") + if total > self.graph_num_tokens: + raise ValueError( + f"packed tokens {total} exceed the captured bucket {self.graph_num_tokens}" + ) + if exact_fill and total != self.graph_num_tokens: + raise ValueError( + f"packed tokens {total} do not fill the captured bucket " + f"{self.graph_num_tokens}; seq_lens.sum() would disagree with " + f"the row count the graph was captured for" + ) + expected = build_qo_indptr(self.verify_lens) + if not torch.equal(self.qo_indptr, expected): + raise ValueError("qo_indptr is not the prefix sum of verify_lens") + + +@dataclass(frozen=True) +class RaggedCaptureShape: + """The ``(rows, tokens)`` pair a ragged batch is padded up to. + + Attributes: + padded_bs: captured batch size (number of rows, real + pad requests). + bucket: captured token count, i.e. ``sum(verify_lens)`` after padding. + """ + + padded_bs: int + bucket: int + + +def _bucket_decomposable(bucket: int, rows: int, tokens: int, padded_bs: int, + max_verify_len: int) -> bool: + """Can a rank with ``rows`` real requests (``tokens`` floor) realise ``bucket``? + + Mirrors the pad-row arithmetic in ``fit_ragged_verify_lens`` exactly: the + ``padded_bs - rows`` pad rows all carry ONE shared window in + ``[1, max_verify_len]`` (they are a single shared dummy object), and the + real rows absorb the remainder within ``[tokens, rows * max_verify_len]``. + + This is computed here, from the allgathered ``(rows, tokens)``, so that + every rank evaluates every OTHER rank's feasibility too: a bucket that + rank 3 cannot decompose used to make rank 3 alone decline the fit while + its peers published the bucket -- a graph-key divergence the ADP shape + gate then turned into a whole-group eager step. Measured on job 2586075 + (pinned full block, where floor == capacity and the divisibility trap is + sharpest): 70 of 962 steps went eager on exactly this. + """ + n_pad = padded_bs - rows + floor, cap = tokens, rows * max_verify_len + if n_pad == 0: + return floor <= bucket <= cap + lo = max(1, -(-(bucket - cap) // n_pad)) + hi = min(max_verify_len, (bucket - floor) // n_pad) + return lo <= hi + + +def choose_ragged_capture_shape( + *, + num_real_requests: int, + total_verify_tokens: int, + bs_buckets: Sequence[int], + token_buckets: Sequence[int], + peer_stats: Optional[Sequence[Sequence[int]]] = None, + max_verify_len: Optional[int] = None, +) -> RaggedCaptureShape: + """Pick the captured ``(padded_bs, bucket)`` for a ragged batch. + + Under attention DP every rank must land on the *same* shape or the ranks + replay different graphs and their collectives diverge, so agreement comes + from rounding a reduction: pass every rank's + ``(num_real_requests, total_verify_tokens)`` as ``peer_stats`` (including + this rank's own) and each rank computes an identical + + padded_bs = round_up_bs(max_r num_real_r) + bucket = round_up_bucket(padded_bs + max_r (total_r - num_real_r)) + + where ``slack_r = total_r - num_real_r`` is the drafted-position count on + rank ``r`` and the ``padded_bs +`` term is the one-token-per-row floor. + Raises when no captured bucket fits; the caller must shrink the batch. + + With ``max_verify_len`` given, the answer is additionally required to be + DECOMPOSABLE BY EVERY RANK (see ``_bucket_decomposable``): the smallest + such bucket is chosen, walking past candidates any rank cannot realise, + so accept/decline is a pure function of the allgathered payload -- all + ranks pick the same bucket or all raise together. + """ + if not bs_buckets: + raise ValueError("choose_ragged_capture_shape requires bs_buckets") + if num_real_requests < 0 or total_verify_tokens < num_real_requests: + raise ValueError( + f"total_verify_tokens {total_verify_tokens} cannot be below " + f"num_real_requests {num_real_requests} (every request verifies " + f"at least one position)" + ) + + stats = list(peer_stats) if peer_stats else [(num_real_requests, total_verify_tokens)] + # Peer entries may carry extra trailing fields; only the first two are read. + max_real = max(int(peer[0]) for peer in stats) + max_slack = max(int(peer[1]) - int(peer[0]) for peer in stats) + + sorted_bs = sorted({int(b) for b in bs_buckets}) + if max_real > sorted_bs[-1]: + raise ValueError( + f"{max_real} requests exceed the largest captured batch size " + f"{sorted_bs[-1]}; the caller must shrink the batch" + ) + padded_bs = sorted_bs[bisect.bisect_left(sorted_bs, max_real)] + + needed = padded_bs + max_slack + candidates = sorted({int(t) for t in token_buckets}) + if max_verify_len is None: + bucket = round_up_to_bucket(needed, candidates) + return RaggedCaptureShape(padded_bs=padded_bs, bucket=bucket) + # Group-consistent choice: smallest captured bucket every rank can + # decompose. The floor requirement (bucket >= tokens_r + pad rows at one + # token each, i.e. >= needed at the widest rank) is subsumed by the + # decomposability test, so no separate >= needed cut is applied. + for cand in candidates: + if all( + _bucket_decomposable(cand, int(peer[0]), int(peer[1]), + padded_bs, int(max_verify_len)) + for peer in stats): + return RaggedCaptureShape(padded_bs=padded_bs, bucket=cand) + raise ValueError( + f"no captured bucket (grid {candidates}) is decomposable by every " + f"rank (peer stats {list(stats)}, padded_bs {padded_bs}, " + f"max_verify_len {max_verify_len}); the group falls back to uniform " + f"together rather than splitting its graph keys") + + +def fill_bucket_device( + verify_lens: torch.Tensor, + *, + num_real: torch.Tensor, + graph_num_tokens: int, + max_verify_len: int, + pad_fill: Optional[int] = None, +) -> torch.Tensor: + """Capture-safe :meth:`RaggedVerifyLayout.fill_bucket`: same allocation, + computed entirely on device. + + The host round-robin decomposes into full cycles -- every row with + headroom gains one token -- plus one final partial cycle where the first + ``spare`` headroom rows in index order gain one. Each cycle is therefore + ``grant = headroom & (cumsum(headroom) <= spare)``, and a row's headroom + is at most ``max_verify_len - 1``, so a statically-unrolled loop of that + many cycles per phase (real rows first, then pad rows) reproduces the + host result token-for-token with no ``.item()`` or host sync. + + The device version cannot raise, so feasibility is the caller's contract, + checkable from host-known quantities alone before launch: + ``padded_bs <= graph_num_tokens <= padded_bs * max_verify_len`` and the + scheduled budget small enough that real tokens plus one-per-pad-row fit + the bucket. Under those bounds the result sums to exactly + ``graph_num_tokens``. + + Args: + verify_lens: ``[padded_bs]`` scheduled lengths; entries at or beyond + ``num_real`` are ignored (pad rows restart from one token). + num_real: 0-d integer tensor on the same device (host-known per step, + but not a capture constant -- it changes across replays). + graph_num_tokens: the captured token bucket (capture constant). + max_verify_len: per-row ceiling (capture constant). + pad_fill: when given, every pad row carries EXACTLY this many tokens + and all slack goes to real rows. The host fit publishes this split + (``ragged_pad_verify_len``) because its copy widths -- how many + flat tokens belong to real rows -- must be known before launch; + the device fill has to land on the same split or the host and + device disagree about where real tokens end. + Returns: + ``[padded_bs]`` int32 filled lengths. + """ + if verify_lens.dim() != 1: + raise ValueError(f"verify_lens must be 1-D, got {tuple(verify_lens.shape)}") + padded_bs = verify_lens.numel() + device = verify_lens.device + rows = torch.arange(padded_bs, device=device) + is_real = rows < num_real + lens = torch.where( + is_real, + verify_lens.to(torch.int32), + torch.full((padded_bs,), 1 if pad_fill is None else int(pad_fill), + dtype=torch.int32, device=device), + ) + spare = graph_num_tokens - lens.sum() + phases = (is_real,) if pad_fill is not None else (is_real, ~is_real) + for phase_mask in phases: + for _ in range(max(max_verify_len - 1, 0)): + headroom = phase_mask & (lens < max_verify_len) + in_cycle = torch.cumsum(headroom.to(torch.int64), dim=0) <= spare + grant = (headroom & in_cycle).to(torch.int32) + lens = lens + grant + spare = spare - grant.sum() + return lens + + +def build_row_maps_device( + verify_lens: torch.Tensor, + *, + graph_num_tokens: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-token ``(req_idx, kv_correction)`` from filled lengths, on device. + + The token-major row maps prepare() stages from host lists decompose into + a pure function of the filled ``verify_lens``: token ``t`` of a request + with window ``v`` at within-request offset ``o`` attends KV extent + ``kv_len - v + 1 + o``, i.e. correction ``o - v + 1``. Composing the + returned correction with a ``kv_lens`` gather is exactly + ``refresh_ragged_row_kv_lens``, so the caller reuses that op unchanged. + + ``verify_lens`` must already sum to ``graph_num_tokens`` + (post-:func:`fill_bucket_device`); the explicit total keeps + ``repeat_interleave`` sync-free. + + Returns: + ``req_idx``: ``[graph_num_tokens]`` int64 owning-row index. + ``correction``: ``[graph_num_tokens]`` int32, in ``[1 - v, 0]``. + """ + indptr = build_qo_indptr(verify_lens).to(torch.long) + req_idx = row_ids_from_lens(verify_lens, total=graph_num_tokens) + offset = torch.arange(graph_num_tokens, device=verify_lens.device) - indptr[req_idx] + correction = (offset - verify_lens.to(torch.long)[req_idx] + 1).to(torch.int32) + return req_idx, correction + + +def ragged_gather_index_lists( + slots: Sequence[int], counts: Sequence[int] +) -> Tuple[List[int], List[int]]: + """Row/column index lists for gathering a ragged block from a + ``[num_slots, max_width]`` tensor. + + ``slots[i]`` contributes ``counts[i]`` entries from columns + ``0..counts[i]-1``, concatenated in batch order; with a constant ``counts`` + this matches the uniform ``tensor[slots, :width]`` gather. + """ + if len(slots) != len(counts): + raise ValueError( + f"slots and counts must be the same length, got {len(slots)} and {len(counts)}" + ) + rows: List[int] = [] + cols: List[int] = [] + for slot, count in zip(slots, counts): + if count < 0: + raise ValueError(f"negative gather count {count} for slot {slot}") + rows.extend([slot] * count) + cols.extend(range(count)) + return rows, cols + + +def row_ids_from_lens(verify_lens: torch.Tensor, *, total: int) -> torch.Tensor: + """``[bs]`` lengths -> ``[total]`` owning-request id for each packed token. + + ``total`` (``sum(verify_lens)``) is mandatory: without ``output_size``, + ``repeat_interleave`` with device-resident repeats syncs to size its + output, which is illegal inside the captured graph. + """ + return torch.repeat_interleave( + torch.arange(verify_lens.numel(), device=verify_lens.device), + verify_lens.to(torch.long), + output_size=int(total), + ) + + +def scatter_ragged_to_padded( + flat: torch.Tensor, + *, + verify_lens: torch.Tensor, + qo_indptr: torch.Tensor, + max_len: int, + pad_value: int = 0, +) -> torch.Tensor: + """Unpack a flat ragged batch into ``[bs, max_len, ...]`` with padding. + + Sync-free (capturable): token ``t`` goes to request ``row_ids[t]`` at + column ``t - qo_indptr[row_ids[t]]``, and ``flat.shape[0]`` supplies the + ``output_size`` :func:`row_ids_from_lens` needs. + """ + bs = int(verify_lens.numel()) + rows = row_ids_from_lens(verify_lens, total=flat.shape[0]) + cols = torch.arange(flat.shape[0], device=flat.device) - qo_indptr.to(torch.long)[rows] + out = flat.new_full((bs, max_len, *flat.shape[1:]), pad_value) + out[rows, cols] = flat + return out + + +def fill_padded_rows_onehot( + probs: torch.Tensor, + *, + verify_lens: torch.Tensor, +) -> torch.Tensor: + """Make ``[bs, max_len, vocab]`` padding rows a valid one-hot distribution. + + Padding rows are never read, but a rejection-sampling kernel still walks + them and sampling an all-zero row can index out of bounds; all mass goes + to token 0, which affects no position that is read. Modifies ``probs`` in + place and returns it. + """ + positions = torch.arange(probs.shape[1], device=probs.device) + invalid = positions.unsqueeze(0) >= verify_lens.to(positions.dtype).unsqueeze(1) + probs[..., 0] = torch.where(invalid, torch.ones_like(probs[..., 0]), probs[..., 0]) + return probs + + +def count_accepted_ragged( + *, + draft_tokens: torch.Tensor, + target_tokens: torch.Tensor, + verify_lens: torch.Tensor, +) -> torch.Tensor: + """Accepted draft-token count per request, honouring per-request windows. + + Positions at or beyond a request's ``verify_len`` hold stale padding and + must be masked to "no match" before the cumulative product, or a stale + slot can credit an acceptance the target never made. + + Args: + draft_tokens: ``[bs, max_len]`` padded drafted tokens. + target_tokens: ``[bs, max_len]`` padded target tokens. + verify_lens: ``[bs]`` positions actually verified per request. + """ + if draft_tokens.shape != target_tokens.shape: + raise ValueError( + f"draft {tuple(draft_tokens.shape)} and target " + f"{tuple(target_tokens.shape)} must have the same padded shape" + ) + max_len = draft_tokens.shape[1] + positions = torch.arange(max_len, device=draft_tokens.device) + valid = positions.unsqueeze(0) < verify_lens.to(positions.dtype).unsqueeze(1) + match = (draft_tokens == target_tokens) & valid + return torch.cumprod(match.int(), dim=-1).sum(1) diff --git a/tensorrt_llm/_torch/speculative/dspark_schedule.py b/tensorrt_llm/_torch/speculative/dspark_schedule.py new file mode 100644 index 000000000000..643585b32e63 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_schedule.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DSpark confidence-scheduled verification (arXiv:2607.05147 §5). + +``survival[r][j] = prod(conf[r][:j+1])`` ranks every ``(request, position)`` +candidate globally; the best ``budget`` are admitted. Survival is +non-increasing along each request's positions, so any global top-k is +automatically a per-request prefix and the allocation only counts admitted +candidates per row. Only the number of verified tokens is decided here, never +acceptance, so scheduling is lossless w.r.t. the target distribution. Ties are +broken by ``(position, request)``, never by value: ranks with +bitwise-different confidences must still choose identical lengths or their +batch shapes (and collectives) diverge. +""" + +from dataclasses import dataclass +from typing import Union + +import torch + +__all__ = [ + "DSparkScheduleConfig", + "compute_survival", + "schedule_verify_lens_topk", +] + + +# Raw logit for confidence rows that carry no measurement: sigmoid(30/T) ~ 1.0 +# at any fitted temperature, so an unknown row is treated as "certainly accept" +# -- optimistic by design. Lives here because the worker (writes it) and the +# planner (counts it) cannot import each other. +NEUTRAL_CONFIDENCE_LOGIT = 30.0 + + +@dataclass(frozen=True) +class DSparkScheduleConfig: + """Bounds for the per-request verify length. + + Attributes: + block_size: draft block length ``K``. + min_verify_len: floor on every request's verify length; must stay >= 1 + (position 0 carries the bonus/anchor token). Budget is allocated + *above* this floor. + max_verify_len: cap on every request's verify length; 0 means + ``block_size``; never exceeds ``block_size``. + survival_eps: numerical floor (not a tuning knob); candidates below it + are dropped before ranking. + """ + + block_size: int + min_verify_len: int = 1 + max_verify_len: int = 0 + survival_eps: float = 1e-6 + + def __post_init__(self) -> None: + if self.block_size < 1: + raise ValueError(f"block_size must be >= 1, got {self.block_size}") + if self.min_verify_len < 1: + raise ValueError( + f"min_verify_len must be >= 1 (every request verifies its anchor " + f"token), got {self.min_verify_len}" + ) + if self.resolved_max_verify_len < self.min_verify_len: + raise ValueError( + f"max_verify_len ({self.resolved_max_verify_len}) < min_verify_len " + f"({self.min_verify_len})" + ) + if not 0.0 <= self.survival_eps < 1.0: + raise ValueError(f"survival_eps must be in [0, 1), got {self.survival_eps}") + + @property + def resolved_max_verify_len(self) -> int: + """``max_verify_len`` with 0 meaning "the whole block", clamped to ``block_size``.""" + cap = self.max_verify_len or self.block_size + return min(int(cap), int(self.block_size)) + + @property + def schedulable_per_request(self) -> int: + """Positions a single request can win from the budget, above the floor.""" + return self.resolved_max_verify_len - self.min_verify_len + + +def compute_survival(confidence: torch.Tensor) -> torch.Tensor: + """``[bs, K]`` conditional acceptance probabilities -> prefix survival. + + ``out[r][j] = prod_{i <= j} confidence[r][i]``, computed in fp32. The result + is non-increasing along ``j``, which :func:`schedule_verify_lens_topk` relies + on to keep its allocation prefix-shaped. + """ + if confidence.dim() != 2: + raise ValueError(f"confidence must be [bs, K], got shape {tuple(confidence.shape)}") + return torch.cumprod(confidence.to(torch.float32), dim=1) + + +def schedule_verify_lens_topk( + *, + survival: torch.Tensor, + budget: Union[int, torch.Tensor], + cfg: DSparkScheduleConfig, +) -> torch.Tensor: + """Allocate ``budget`` verify tokens across the batch by global survival rank. + + Args: + survival: ``[bs, K]`` prefix-survival probabilities from + :func:`compute_survival`. + budget: total verify tokens to hand out *above* the per-request floor. + A 0-d integer tensor is honoured without a host sync (the value + stays device-resident through the ``rank < budget`` cut), which is + what lets the whole function run inside a captured graph. + cfg: bounds; see :class:`DSparkScheduleConfig`. + Returns: + ``[bs]`` int32 verify lengths, each in + ``[min_verify_len, resolved_max_verify_len]``, summing to at most + ``bs * min_verify_len + budget``. + + Determinism: ranking uses a value-independent tie-break (position before + request), so equal survivals resolve identically regardless of float noise + or backend. + """ + if survival.dim() != 2: + raise ValueError(f"survival must be [bs, K], got shape {tuple(survival.shape)}") + bs, num_positions = survival.shape + if num_positions != cfg.block_size: + raise ValueError( + f"survival has {num_positions} positions but block_size is {cfg.block_size}" + ) + + device = survival.device + floor = int(cfg.min_verify_len) + schedulable = int(cfg.schedulable_per_request) + verify_lens = torch.full((bs,), floor, dtype=torch.int32, device=device) + if not isinstance(budget, torch.Tensor): + budget = int(budget) + if budget <= 0: + return verify_lens + if bs == 0 or schedulable <= 0: + return verify_lens + + # Positions [0, floor) are already granted to every request; only those + # above compete for budget. + candidates = survival[:, floor : floor + schedulable].to(torch.float32) + flat = candidates.reshape(-1) + eligible = flat >= cfg.survival_eps + + # Sort the tie-break key first, then stable-sort on the scores: the final + # order is a pure function of (survival, position, request). Ineligible + # candidates score -1 and are masked by ``& eligible`` below; no + # ``.item()``/host sync anywhere, so the path stays capture-safe. + scores = torch.where(eligible, flat, torch.full_like(flat, -1.0)) + tie_break = torch.arange(schedulable, device=device).repeat(bs) * bs + torch.arange( + bs, device=device + ).repeat_interleave(schedulable) + order = torch.argsort(tie_break, stable=True) + order = order[torch.argsort(scores[order], descending=True, stable=True)] + + rank = torch.empty_like(order) + rank[order] = torch.arange(order.numel(), device=device) + selected = (rank < budget) & eligible + + # survival is non-increasing along positions, so the chosen set for a request + # is necessarily a prefix -- counting per row recovers the full allocation. + granted = selected.view(bs, schedulable).sum(dim=1).to(torch.int32) + verify_lens += granted + return torch.clamp(verify_lens, min=floor, max=cfg.resolved_max_verify_len) diff --git a/tensorrt_llm/_torch/speculative/dspark_sts.py b/tensorrt_llm/_torch/speculative/dspark_sts.py new file mode 100644 index 000000000000..40f3d50afc2b --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_sts.py @@ -0,0 +1,295 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""STS calibration: collecting the data, and reading the fitted table. + +``apply_sts`` computes ``confidence[r][j] = sigmoid(logit[r][j] / T[j])``. The +planner consumes the cumulative product ``survival[r][j]``, in which +per-position calibration error compounds geometrically, so the table must be +fitted against the survival (see ``tests/microbenchmarks/dspark_fit_sts.py``). +Collection is opt-in via ``TLLM_DSPARK_STS_COLLECT_PATH`` and is not a serving +path: it reads the confidence buffer back to the host once per decode step. +""" + +import json +import os +from pathlib import Path +from typing import List, Optional + +import torch + +from ..._utils import prefer_pinned +from ...logger import logger + +__all__ = [ + "STS_COLLECT_ENV", + "DSparkStsRecorder", + "load_sts_temperatures_from_path", + "make_recorder_from_env", +] + +#: Set to a path stem to collect calibration data. Shards land at +#: ``...pt``. +STS_COLLECT_ENV = "TLLM_DSPARK_STS_COLLECT_PATH" + +#: Flush after this many recorded steps, so a killed run still leaves data. +_DEFAULT_FLUSH_EVERY = 64 + + +def load_sts_temperatures_from_path(path: str) -> List[float]: + """Read a temperature vector, accepting either spelling of the key. + + This repo writes ``sts_temperatures``; SGLang writes ``temperatures``. + The vectors are interchangeable. + """ + with open(path, encoding="utf-8") as handle: + payload = json.load(handle) + for key in ("sts_temperatures", "temperatures"): + if key in payload: + temps = payload[key] + break + else: + raise KeyError( + f"{path} has neither 'sts_temperatures' nor 'temperatures'; " + f"found keys {sorted(payload)}") + if not temps: + raise ValueError(f"{path} carries an empty temperature vector") + return [float(t) for t in temps] + + +def resolve_confidence_head(draft_model): + """Find the confidence head across the known draft-model layouts. + + Two layouts exist: the bare ``DSparkDraftModel`` (stages under + ``.mtp_layers``, head on the last stage) and the ``DSparkForCausalLM`` + wrapper (bare model under ``.dspark_model``). Written once and unit-tested + against both; returns None when no head is found. + """ + inner = getattr(draft_model, "dspark_model", draft_model) + stages = getattr(inner, "mtp_layers", None) + if not stages: + return None + return getattr(stages[-1], "confidence_head", None) + + +class DSparkStsRecorder: + """Pair each drafted block's logits with how much of it was accepted. + + The confidence buffer is overwritten by later draft passes before the + accepted count arrives, so the join is made on two axes, never wall-clock + arrival: by draft-pass stamp (``stage_snapshot`` keeps a ring of host + copies keyed by the worker's draft sequence counter) and by buffer row + resolved through the worker's own allocator (``confidence_row_for`` -- + never ``py_seq_slot``, which belongs to a different allocator). A pair + that cannot be made exactly is dropped and counted in ``stats``, never + approximated. + """ + + #: Pairing lag is one draft pass (two under the overlap scheduler); four + #: clears both without letting a wrong target_seq match ancient content. + _RING_DEPTH = 4 + + def __init__(self, *, path_stem: str, block_size: int, rank: int = 0, + flush_every: int = _DEFAULT_FLUSH_EVERY) -> None: + self.path_stem = str(path_stem) + self.block_size = int(block_size) + self.rank = int(rank) + self.flush_every = max(1, int(flush_every)) + self._logits: List[torch.Tensor] = [] + self._accepted: List[int] = [] + self._shard = 0 + self._steps_since_flush = 0 + # Snapshot ring, indexed by staged_seq % _RING_DEPTH. Each entry is + # (host_logits, host_stamps, cuda_event, staged_seq). `record` selects + # the entry whose stamp for the request's row equals the pass that + # drafted the block being verified: pairing is by draft-pass identity, + # never wall-clock arrival. + self._ring: List[Optional[tuple]] = [None] * self._RING_DEPTH + # Every decline is counted; a thin shard must explain itself. + self.stats: dict = {"recorded": 0, "no_snapshot": 0, + "stale_stamp": 0, "row_out_of_range": 0, + "no_row": 0, "snapshots_staged": 0} + + def stage_snapshot(self, *, device_logits: torch.Tensor, + device_stamps: Optional[torch.Tensor], + staged_seq: Optional[int]) -> None: + """Snapshot the confidence buffer + stamps, keyed by draft pass. + + Must be called from the executor's host path once per step, never from + inside the captured graph: Python in a captured region runs at capture + time only. + """ + if device_stamps is None or staged_seq is None: + return + idx = int(staged_seq) % self._RING_DEPTH + entry = self._ring[idx] + if (entry is None or entry[0].shape != device_logits.shape + or entry[1].shape != device_stamps.shape): + on_gpu = device_logits.is_cuda + entry = ( + torch.empty(device_logits.shape, dtype=torch.float32, + device="cpu", + pin_memory=on_gpu and prefer_pinned()), + torch.empty(device_stamps.shape, dtype=torch.int32, + device="cpu", + pin_memory=on_gpu and prefer_pinned()), + torch.cuda.Event() if on_gpu else None, + None, + ) + entry[0].copy_(device_logits.detach().to(torch.float32), + non_blocking=True) + entry[1].copy_(device_stamps, non_blocking=True) + if entry[2] is not None: + entry[2].record() + self._ring[idx] = (entry[0], entry[1], entry[2], int(staged_seq)) + self.stats["snapshots_staged"] += 1 + + def record(self, *, row: Optional[int], accepted: int, + target_seq: Optional[int]) -> None: + """Pair one request's accepted count with the logits that drafted it. + + Args: + row: the request's row in the confidence buffer, resolved through + the worker's own allocator (``confidence_row_for``), never + ``py_seq_slot``, which belongs to a different allocator. + accepted: drafted positions accepted, excluding the bonus token. + target_seq: the draft pass that produced the block this label + verifies, snapshotted into the SampleState at sampling time + (reading it live rewinds wrong under the overlap scheduler). + """ + if row is None: + self.stats["no_row"] += 1 + return + if target_seq is None: + self.stats["no_snapshot"] += 1 + return + entry = self._ring[int(target_seq) % self._RING_DEPTH] + if entry is None: + self.stats["no_snapshot"] += 1 + return + logits_host, stamps_host, event, staged_seq = entry + if not 0 <= int(row) < logits_host.shape[0]: + self.stats["row_out_of_range"] += 1 + return + if event is not None: + event.synchronize() + if int(stamps_host[int(row)]) != int(target_seq): + # This row's content is from a different draft pass than the label + # verifies; refusing is the point -- appending would mispair. + self.stats["stale_stamp"] += 1 + return + self._logits.append(logits_host[int(row)].clone()) + self._accepted.append(int(accepted)) + self.stats["recorded"] += 1 + + def end_step(self) -> None: + self._steps_since_flush += 1 + if self._steps_since_flush >= self.flush_every: + self.flush() + + def flush(self) -> None: + """Write a shard. Cheap to call with nothing buffered.""" + self._steps_since_flush = 0 + if not self._logits: + return + logits = torch.stack(self._logits, dim=0) + counts = torch.tensor(self._accepted, dtype=torch.int64).view(-1, 1) + positions = torch.arange(self.block_size).view(1, -1) + # prefix_mask[i, k] == 1 iff draft positions 0..k were ALL accepted, + # which is the event whose probability `survival[k]` estimates. Any + # other label (e.g. "position k accepted") would fit a different + # quantity than the planner consumes. + prefix_mask = (positions < counts).to(torch.float32) + + path = Path(f"{self.path_stem}.r{self.rank}.{self._shard}.pt") + path.parent.mkdir(parents=True, exist_ok=True) + # The pairing marker lets the fitter reject mispaired legacy shards; + # the decline counters ride along so a thin shard explains itself. + payload = { + "logits": logits, + "prefix_mask": prefix_mask, + "meta": { + "pairing": "draft_seq_ring", + "block_size": self.block_size, + "stats": dict(self.stats), + }, + } + torch.save(payload, path) + logger.info( + f"DSpark STS: wrote {logits.shape[0]} samples to {path} " + f"(stats={self.stats})") + self._logits.clear() + self._accepted.clear() + self._shard += 1 + + +def make_recorder_from_env(*, block_size: int, rank: int = 0, + has_cost_table: bool = False, + ragged_mode: Optional[str] = None + ) -> Optional[DSparkStsRecorder]: + """Build a recorder iff ``TLLM_DSPARK_STS_COLLECT_PATH`` is set. + + Refuses the two regimes in which the collected pairs describe the + scheduler rather than the head; both would otherwise fail silently. + """ + stem = os.environ.get(STS_COLLECT_ENV, "").strip() + if not stem: + return None + + forced = os.environ.get("TLLM_DSPARK_FORCE_VERIFY_LEN", "").strip() + forced_full = False + if forced: + try: + forced_full = int(forced) >= int(block_size) + except ValueError: + forced_full = False + if has_cost_table and not forced_full: + raise ValueError( + f"{STS_COLLECT_ENV} is set while a cost table is loaded. The planner " + f"will trim, and a trimmed window censors the acceptance label -- the " + f"fit would describe the scheduler, not the confidence head. Either " + f"collect without a cost table, or pin the window with " + f"TLLM_DSPARK_FORCE_VERIFY_LEN= (a pin below the " + f"block still censors the label and is refused).") + + if (ragged_mode or "").lower() == "compact": + raise ValueError( + f"{STS_COLLECT_ENV} is set with ragged verify mode 'compact'. Padded " + f"verify rows corrupt the per-position prefix label. Collect in " + f"static mode (TLLM_DSPARK_RAGGED_VERIFY_MODE=static).") + + if os.environ.get("TLLM_DSPARK_DEVICE_WINDOWS", "0") == "1": + raise ValueError( + f"{STS_COLLECT_ENV} is set with TLLM_DSPARK_DEVICE_WINDOWS=1. " + f"Device-selected windows trim per request on device, so every " + f"collected acceptance label is censored by a window the host " + f"never sees. Collect with device windows off.") + + frac = os.environ.get("TLLM_DSPARK_FORCE_BUDGET_FRAC", "").strip() + if frac: + try: + frac_full = float(frac) >= 1.0 + except ValueError: + frac_full = False + if not frac_full: + raise ValueError( + f"{STS_COLLECT_ENV} is set while TLLM_DSPARK_FORCE_BUDGET_FRAC=" + f"{frac!r} trims the verify budget; a trimmed window censors " + f"the acceptance label for every collected sample. Collect " + f"with the fraction unset or at 1.0.") + logger.warning( + f"DSpark STS collection ON ({STS_COLLECT_ENV}={stem!r}). This adds one " + f"device->host copy per decode step; it is a calibration mode, not a " + f"serving configuration.") + return DSparkStsRecorder(path_stem=stem, block_size=block_size, rank=rank) diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py new file mode 100644 index 000000000000..d3b530cbb3fc --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -0,0 +1,653 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Host-side coordinator turning DSpark confidence into a per-iteration draft length. + +Contracts: the decision is made on the host (CUDA-graph selection is a host +value); the verified token total must round up to a captured verify bucket, or +the step silently runs eager; and every rank must shape the batch identically +-- the caller reduces the ragged decision across ranks before any rank acts on +it, while :meth:`decide_verify_lens` itself is rank-local by default. +Confidence is read from a lagged, non-blocking snapshot; a snapshot that is +not ready degrades to verifying everything, never to a stale guess. +""" + +import os +from typing import Callable, List, Optional, Sequence, Tuple + +import numpy as np +import torch + +from ..._utils import prefer_pinned +from ...logger import logger +from .dspark_planner import SpsCostTable, compute_verify_token_budget +from .dspark_schedule import (NEUTRAL_CONFIDENCE_LOGIT, DSparkScheduleConfig, + compute_survival, schedule_verify_lens_topk) + +#: Profiling-only pin (SPS profiling / STS collection): fixes every request's +#: verify window at this length, bypassing the planner. Not a serving knob. +FORCE_VERIFY_LEN_ENV = "TLLM_DSPARK_FORCE_VERIFY_LEN" + +#: Profiling-only budget override (SPS profiling on the live ragged path): +#: replaces the cost-table argmax with ``frac`` of the maximum trimmable +#: budget, while windows still flow through the real confidence top-k. Unlike +#: the verify-length pin this keeps the executed shape production-faithful +#: (ragged windows, bucket fit); the profiler sweeps it to visit off-diagonal +#: (bs, M) cells. Not a serving knob. +FORCE_BUDGET_FRAC_ENV = "TLLM_DSPARK_FORCE_BUDGET_FRAC" + +#: Wire encoding for the frac's cross-rank agreement: fractions travel as +#: fixed-point ints so the allgather payload stays homogeneous ints. +_FRAC_WIRE_SCALE = 10_000 + +#: Device-side window selection (the paper's dual-timescale split): the host +#: decides only the CAPACITY -- (padded_bs, bucket, budget) -- from a snapshot +#: one step older than today's (lag-2 by the block-relative count, matching +#: SGLang), and a pre-replay device prologue re-ranks the batch with the +#: verified block's own fresh confidence (lag-0). Constant per process, read +#: once at init on every rank. +DEVICE_WINDOWS_ENV = "TLLM_DSPARK_DEVICE_WINDOWS" + +#: Experiment-only mixed-step guard: when set, a rank whose scheduled batch +#: carries context (prefill) requests declines to trim that step, and the +#: cross-rank all-or-nothing gate reverts the whole group to the full block. +#: The cost table is profiled on pure-decode steps and cannot price a trim +#: whose step time is dominated by prefill compute; this isolates that effect. +SKIP_MIXED_TRIM_ENV = "TLLM_DSPARK_SKIP_MIXED_TRIM" + + +def device_windows_enabled() -> bool: + """Whether device-side window selection is on (process-constant env).""" + return os.environ.get(DEVICE_WINDOWS_ENV, "0") == "1" + + +__all__ = ["DSparkVerifyPlanner"] + + +class DSparkVerifyPlanner: + """Stages confidence off the device and turns it into a draft length. + + Args: + cfg: verify-length bounds. + cost_table: measured step cost. A flat table means "not profiled"; the + planner then refuses to trim (see :meth:`decide_verify_lens`). + tiers: draft lengths that have a captured CUDA graph. The drafted block + is always ``max_tier``; these also bound the budget search. + apply_calibration: maps raw confidence logits to probabilities. Normally + ``confidence_head.apply_sts``; defaults to a plain sigmoid. + all_rank_max: cross-rank reduction, returning the max of ``value`` over + ranks. ``None`` for single-rank runs. + """ + + def __init__( + self, + *, + cfg: DSparkScheduleConfig, + cost_table: Optional[SpsCostTable] = None, + tiers: Optional[Sequence[int]] = None, + apply_calibration: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + all_rank_max: Optional[Callable[[int], int]] = None, + ): + self.cfg = cfg + self.cost_table = cost_table + self.tiers: List[int] = sorted({int(t) for t in (tiers or [cfg.resolved_max_verify_len])}) + self.apply_calibration = apply_calibration or torch.sigmoid + self.all_rank_max = all_rank_max + + # Profiling pin, read once at init: constant per process and forwarded + # to every worker at spawn, so all ranks agree without a collective. + # An absolute length (not a fraction) so M = bs * (L + 1) lands exactly + # on a captured bucket. + self._forced_verify_len = self._read_forced_verify_len() + # Runtime pin queued but not yet in force, held as its wire value: + # -1 nothing queued, 0 clear the pin, >0 pin to that length + # (min_verify_len >= 1, so 0 is unambiguous). Adopted through the + # step's cross-rank allgather so every rank starts pinning on the + # same step; applying it locally would diverge the shape gate. + self._pending_verify_len_pin: int = -1 + # Budget-fraction override (same lifecycle as the pin: env at init, + # runtime requests adopted through the step's allgather). Wire values: + # -1 nothing queued, 0 clear, >0 frac * _FRAC_WIRE_SCALE. + self._forced_budget_frac = self._read_forced_budget_frac() + self._pending_budget_frac: int = -1 + + self._host_buffer: Optional[torch.Tensor] = None + self._copy_event: Optional[torch.cuda.Event] = None + self._host_stamps: Optional[torch.Tensor] = None + self._staged_seq: Optional[int] = None + self._snapshot_valid = False + + # Device-window mode: ranking happens on device with fresh confidence, + # so the host snapshot serves ONLY the capacity argmax -- and that one + # deliberately reads a snapshot one staging older (lag-2), per the + # paper's design: capacity tolerates staleness, ranking does not. + # Two pinned buffers alternate; the pair we are NOT writing this step + # is the older, guaranteed-landed one. + self.device_windows = os.environ.get(DEVICE_WINDOWS_ENV, "0") == "1" + # Constant per process like the pin: every rank reads the same env at + # spawn, so the per-step decline needs no extra collective. + self.skip_mixed_trim = os.environ.get(SKIP_MIXED_TRIM_ENV, "0") == "1" + if self.device_windows and self._forced_verify_len is not None: + raise ValueError( + f"{DEVICE_WINDOWS_ENV}=1 is incompatible with " + f"{FORCE_VERIFY_LEN_ENV}: the pin fixes host windows that " + f"device selection would silently override") + self._prev_buffer: Optional[torch.Tensor] = None + self._prev_event: Optional[torch.cuda.Event] = None + self._prev_stamps: Optional[torch.Tensor] = None + self._prev_seq: Optional[int] = None + self._prev_valid = False + # Count every path that declines to trim; degradation here is silent. + self.stats = { + "decisions": 0, + "fallback_no_snapshot": 0, + "fallback_flat_cost": 0, + "fallback_no_confidence": 0, + "fallback_short_snapshot": 0, + "fallback_no_gen_requests": 0, + "forced_steps": 0, + } + + @property + def max_tier(self) -> int: + return self.tiers[-1] if self.tiers else self.cfg.resolved_max_verify_len + + def stage_confidence(self, confidence_logits: torch.Tensor, + stamps: Optional[torch.Tensor] = None, + staged_seq: Optional[int] = None) -> None: + """Start a non-blocking device->host copy of the confidence buffer. + + ``confidence_logits`` is the worker's whole slot-indexed buffer, staged + in full so :meth:`decide_verify_lens` can look up any live request by + slot across the one-iteration lag. Never synchronizes: the copy is + issued on the current stream and an event is recorded; + :meth:`decide_verify_lens` polls the event and skips a snapshot that + has not landed. + """ + if confidence_logits is None or confidence_logits.shape[0] == 0: + return + if self.device_windows: + # Rotate before overwriting: what was current becomes the older + # (lag-2) snapshot the budget argmax reads; the old older buffer's + # storage is reused for this staging. Pure reference swaps. + self._host_buffer, self._prev_buffer = (self._prev_buffer, + self._host_buffer) + self._copy_event, self._prev_event = (self._prev_event, + self._copy_event) + self._host_stamps, self._prev_stamps = (self._prev_stamps, + self._host_stamps) + self._staged_seq, self._prev_seq = self._prev_seq, self._staged_seq + self._prev_valid = self._snapshot_valid + self._snapshot_valid = False + if self._host_buffer is None or self._host_buffer.shape != confidence_logits.shape: + self._host_buffer = torch.empty( + confidence_logits.shape, + dtype=torch.float32, + device="cpu", + pin_memory=prefer_pinned(), + ) + self._copy_event = torch.cuda.Event() + self._host_buffer.copy_(confidence_logits, non_blocking=True) + if stamps is not None: + if (self._host_stamps is None + or self._host_stamps.shape != stamps.shape): + self._host_stamps = torch.empty(stamps.shape, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned()) + self._host_stamps.copy_(stamps, non_blocking=True) + self._staged_seq = staged_seq + self._copy_event.record() + self._snapshot_valid = True + + def _ready_snapshot(self) -> Optional[torch.Tensor]: + if not self._snapshot_valid or self._host_buffer is None: + return None + # query() is non-blocking; a not-yet-landed copy just means this step + # decides without it. + if self._copy_event is not None and not self._copy_event.query(): + return None + return self._host_buffer + + _FORCE_ENV = FORCE_VERIFY_LEN_ENV + + def _read_forced_verify_len(self) -> Optional[int]: + """Parse and validate the profiling pin, or None.""" + raw = os.environ.get(self._FORCE_ENV, "").strip() + if not raw: + return None + try: + value = int(raw) + except ValueError: + raise ValueError( + f"{self._FORCE_ENV}={raw!r} is not an integer") from None + lo, hi = int(self.cfg.min_verify_len), int(self.cfg.resolved_max_verify_len) + if not lo <= value <= hi: + raise ValueError( + f"{self._FORCE_ENV}={value} outside [{lo}, {hi}]") + if value not in self.tiers: + raise ValueError( + f"{self._FORCE_ENV}={value} is not in the captured tier ladder " + f"{self.tiers}; every step would fall out of CUDA-graph replay " + f"and the measured times would not be step costs") + logger.warning( + f"DSpark: verify length PINNED to {value} by {self._FORCE_ENV}. " + f"This bypasses the planner entirely and is a profiling mode, not a " + f"serving configuration.") + return value + + def validate_verify_len_pin(self, value: Optional[int]) -> Optional[int]: + """Check a pin against this planner's ladder; None clears the pin.""" + if value is None: + return None + value = int(value) + lo, hi = int(self.cfg.min_verify_len), int(self.cfg.resolved_max_verify_len) + if not lo <= value <= hi: + raise ValueError(f"verify-length pin {value} outside [{lo}, {hi}]") + if value not in self.tiers: + raise ValueError( + f"verify-length pin {value} is not in the captured tier ladder " + f"{self.tiers}; every step would fall out of CUDA-graph replay " + f"and the measured times would not be step costs") + return value + + def _read_forced_budget_frac(self) -> Optional[float]: + raw = os.environ.get(FORCE_BUDGET_FRAC_ENV, "").strip() + if not raw: + return None + try: + value = float(raw) + except ValueError: + raise ValueError( + f"{FORCE_BUDGET_FRAC_ENV}={raw!r} is not a float") from None + value = self.validate_budget_frac(value) + logger.warning( + f"DSpark: verify budget FORCED to {value} of the trimmable " + f"maximum by {FORCE_BUDGET_FRAC_ENV}. Windows still follow " + f"confidence; this is a profiling mode, not a serving " + f"configuration.") + return value + + def validate_budget_frac(self, value: Optional[float]) -> Optional[float]: + """Check a budget fraction; None clears the override.""" + if value is None: + return None + value = float(value) + if not 0.0 < value <= 1.0: + raise ValueError(f"budget fraction {value} outside (0.0, 1.0]") + return value + + @property + def forced_budget_frac(self) -> Optional[float]: + """The budget fraction currently in force, or None.""" + return self._forced_budget_frac + + @property + def forced_verify_len(self) -> Optional[int]: + """The verify-length pin currently in force, or None.""" + return self._forced_verify_len + + def request_budget_frac(self, value: Optional[float]) -> Optional[float]: + """Queue a budget-fraction override (``None`` clears) for the next step. + + Validated here but NOT applied: the step's allgather carries it so all + ranks adopt it together (see :meth:`adopt_budget_frac`). Returns the + wire-quantized value that will actually be adopted, so the caller sees + what applies rather than what it asked for. + """ + value = self.validate_budget_frac(value) + wire = (0 if value is None else + max(1, int(round(value * _FRAC_WIRE_SCALE)))) + self._pending_budget_frac = wire + return None if value is None else min(1.0, wire / _FRAC_WIRE_SCALE) + + def pending_budget_frac(self) -> int: + """The queued request as a wire value: -1 none, 0 clear, >0 fixed-point.""" + return int(self._pending_budget_frac) + + def adopt_budget_frac(self, wire_value: int) -> None: + """Apply the group's agreed budget fraction from the allgathered payload.""" + wire_value = int(wire_value) + if wire_value < 0: + return + # Same non-atomic compare-and-clear as the pin above: an RPC-thread + # request landing between the compare and the store is dropped. Both + # knobs are single-operator profiling controls, so the race is + # tolerated for both rather than fixed for one. + if self._pending_budget_frac == wire_value: + self._pending_budget_frac = -1 + new_frac = (None if wire_value == 0 else + min(1.0, wire_value / _FRAC_WIRE_SCALE)) + if new_frac == self._forced_budget_frac: + return + self._forced_budget_frac = new_frac + logger.warning( + f"DSpark: verify budget fraction set to {new_frac} at runtime. " + f"This bypasses the cost table and is a profiling mode, not a " + f"serving configuration.") + + def request_verify_len_pin(self, value: Optional[int]) -> Optional[int]: + """Queue a pin (or ``None`` to clear) for the next decode step. + + Validated here but NOT applied: the step's allgather carries it so all + ranks adopt it together (see ``adopt_verify_len_pin``). + """ + value = self.validate_verify_len_pin(value) + self._pending_verify_len_pin = 0 if value is None else int(value) + return value + + def pending_verify_len_pin(self) -> int: + """The queued request as a wire value: -1 none, 0 clear, >0 pin.""" + return int(self._pending_verify_len_pin) + + def adopt_verify_len_pin(self, wire_value: int) -> None: + """Apply the group's agreed pin from the allgathered payload. + + Every rank applies the same payload, so the pin takes effect on the + same step everywhere. ``-1`` leaves the current pin alone; ``0`` + clears it. + """ + wire_value = int(wire_value) + if wire_value < 0: + return + # Compare-and-clear: the RPC thread may queue a new request between + # the payload being read and adopted; an unconditional reset would + # drop it. + if self._pending_verify_len_pin == wire_value: + self._pending_verify_len_pin = -1 + new_pin = None if wire_value == 0 else wire_value + if new_pin == self._forced_verify_len: + return + self._forced_verify_len = new_pin + logger.warning( + f"DSpark: verify length pin set to {new_pin} at runtime. This " + f"bypasses the planner and is a profiling mode, not a serving " + f"configuration.") + + def _note_snapshot_stats(self, selected: torch.Tensor, + rows: Optional[Sequence[int]], + stamps: Optional[torch.Tensor] = None, + staged_seq: Optional[int] = None) -> None: + """Record snapshot-quality stats: stamp-lag histogram (staleness per + gathered row) and neutral-row count ("unknown" rows entering the + argmax as certainties). Observability only; never changes behaviour. + ``stamps``/``staged_seq`` override the current snapshot's pair when the + caller read the older (lag-2) one. + """ + try: + stamps = self._host_stamps if stamps is None else stamps + staged_seq = self._staged_seq if staged_seq is None else staged_seq + n = int(selected.shape[0]) + self.stats["snap_rows"] = self.stats.get("snap_rows", 0) + n + neutral = int((selected >= NEUTRAL_CONFIDENCE_LOGIT - 1.0).all(dim=1).sum()) + self.stats["snap_neutral_rows"] = ( + self.stats.get("snap_neutral_rows", 0) + neutral) + if (stamps is not None and staged_seq is not None + and rows is not None): + idx = torch.as_tensor(list(rows), dtype=torch.long) + idx = idx.clamp_(0, stamps.shape[0] - 1) + lags = (int(staged_seq) + - stamps[idx].to(torch.int64)) + hist = self.stats.setdefault("stamp_lag_hist", {}) + for lag, cnt in zip(*[t.tolist() for t in lags.unique( + return_counts=True)]): + key = int(max(-2, min(int(lag), 8))) + hist[key] = hist.get(key, 0) + int(cnt) + except Exception as exc: # noqa: BLE001 - instruments must not kill steps + self.stats["snap_stats_errors"] = ( + self.stats.get("snap_stats_errors", 0) + 1) + if self.stats["snap_stats_errors"] == 1: + logger.warning(f"DSpark snapshot stats failed once: {exc}") + + def _ready_older_snapshot(self) -> Optional[torch.Tensor]: + """The lag-2 snapshot (device-window mode's capacity input), or None. + + By construction its copy was issued a full staging earlier, so its + event has practically always landed; the poll stays for the same + fail-closed reason as :meth:`_ready_snapshot`. + """ + if not self._prev_valid or self._prev_buffer is None: + return None + if self._prev_event is not None and not self._prev_event.query(): + return None + return self._prev_buffer + + def _gather_rows( + self, *, num_gen_requests: int, rows: Optional[Sequence[int]], + snapshot: Optional[torch.Tensor] = None + ) -> Optional[torch.Tensor]: + """This step's confidence, one row per generation request, or None. + + ``rows`` is the per-request buffer row index (``DSparkWorker. + confidence_row_for``); the batch is reshuffled across the one-iteration + lag, so positional reads (``rows=None``) are only correct for callers + that own the ordering. Declines (None) rather than returning fewer + rows than the batch, which would partially assign verify windows + downstream. ``snapshot`` overrides the default (current) snapshot when + the caller already picked one (the lag-2 budget read). + """ + if snapshot is None: + snapshot = self._ready_snapshot() + if snapshot is None: + self.stats["fallback_no_snapshot"] += 1 + return None + if rows is None: + if snapshot.shape[0] < num_gen_requests: + self.stats["fallback_short_snapshot"] += 1 + return None + selected = snapshot[:num_gen_requests] + else: + if len(rows) != num_gen_requests: + self.stats["fallback_short_snapshot"] += 1 + return None + selected = snapshot[torch.as_tensor(list(rows), dtype=torch.long)] + if selected.numel() == 0: + self.stats["fallback_no_confidence"] += 1 + return None + return selected + + def decide_verify_budget( + self, + *, + num_gen_requests: int, + rows: Optional[Sequence[int]] = None, + ) -> Optional[Tuple[int, List[int]]]: + """Capacity only, for device-window mode: ``(budget, shape_lens)``. + + The budget (verify tokens above the per-request floor) comes from the + cost-table argmax over the OLDER staged snapshot -- lag-2 by the + block-relative count, the paper's prescription: capacity tolerates + staleness, ranking does not. ``shape_lens`` is a canonical uniform + spread of the budget used ONLY for the cross-rank shape agreement and + bucket fit (both read totals, never the split); the true per-request + windows are chosen on device from the verified block's own confidence. + Fail-closed like :meth:`decide_verify_lens`: None means verify the + full block. + """ + self.stats["decisions"] += 1 + if num_gen_requests <= 0: + self.stats["fallback_no_gen_requests"] += 1 + return None + n = int(num_gen_requests) + trimmable = self.cfg.resolved_max_verify_len - self.cfg.min_verify_len + if self._forced_budget_frac is not None: + # The frac sweep stays honest in device mode: it replaces only the + # argmax; the device top-k still spends the budget for real. + budget = int(round(self._forced_budget_frac * n * trimmable)) + self.stats["forced_budget_steps"] = ( + self.stats.get("forced_budget_steps", 0) + 1) + else: + if self.cost_table is None or self.cost_table.is_flat: + self.stats["fallback_flat_cost"] += 1 + return None + snapshot = self._ready_older_snapshot() + if snapshot is None: + self.stats["fallback_no_snapshot"] += 1 + return None + selected = self._gather_rows(num_gen_requests=n, rows=rows, + snapshot=snapshot) + if selected is None: + return None + self._note_snapshot_stats(selected, rows, + stamps=self._prev_stamps, + staged_seq=self._prev_seq) + survival = compute_survival(self.apply_calibration(selected)) + budget = compute_verify_token_budget( + survival=survival.numpy().astype(np.float64), + num_gen_requests=n, + cost_table=self.cost_table, + min_verify_len=self.cfg.min_verify_len, + max_verify_len=self.cfg.resolved_max_verify_len, + allowed_lens=self.tiers, + ) + floor_tokens = n * (self.cfg.min_verify_len + 1) + predicted = float(self.cost_table.step_times( + np.asarray([floor_tokens + int(budget)]), n)[0]) + self.last_predicted_step_ms = predicted + self.stats["predicted_ms_sum"] = ( + self.stats.get("predicted_ms_sum", 0.0) + predicted) + self.stats["predicted_steps"] = ( + self.stats.get("predicted_steps", 0) + 1) + budget = max(0, min(int(budget), n * trimmable)) + base, extra = divmod(budget, n) + shape_lens = [ + min(self.cfg.min_verify_len + base + (1 if i < extra else 0), + self.cfg.resolved_max_verify_len) for i in range(n) + ] + return budget, shape_lens + + def decide_verify_lens( + self, + *, + num_gen_requests: int, + rows: Optional[Sequence[int]] = None, + all_rank_max: Optional[Callable[[int], int]] = None, + reduce_across_ranks: bool = True, + budget_override: Optional[int] = None, + ) -> Optional[List[int]]: + """Per-request verify lengths for a ragged step, or None to stay uniform. + + The cost-model argmax sets the budget; the global survival top-k splits + it across requests. Fail-closed: any untrustworthy input (no snapshot, + flat cost table, empty batch) returns None and the step verifies + everything. The batch-wide maximum is reduced across ranks (it sizes + the drafted-token buffer and per-request padding); individual lengths + stay local. Returns exactly ``num_gen_requests`` entries or None, + never a partial list. + """ + # Counted before any decline so the fallback counters have a denominator. + self.stats["decisions"] += 1 + + if num_gen_requests <= 0: + self.stats["fallback_no_gen_requests"] += 1 + return None + + if budget_override is not None: + # Diagnostic: spend a caller-supplied budget (e.g. the lag-2 + # argmax's) through the normal CURRENT-snapshot top-k. Under + # cap-accept this measures a schedule's true acceptance cost -- + # execution stays full-block, the window is only an accounting + # cap -- so two budgets can be compared with identical compute. + selected = self._gather_rows(num_gen_requests=num_gen_requests, + rows=rows) + if selected is None: + return None + self._note_snapshot_stats(selected, rows) + survival = compute_survival(self.apply_calibration(selected)) + lens = schedule_verify_lens_topk(survival=survival, + budget=int(budget_override), + cfg=self.cfg).tolist() + elif self._forced_verify_len is not None: + # Must stay ahead of the cost-table gate (the profiler runs without + # a table) and must set `lens` rather than return, so the pinned + # path still runs the cross-rank agreement below. + self.stats["forced_steps"] += 1 + lens = [int(self._forced_verify_len)] * int(num_gen_requests) + elif self._forced_budget_frac is not None: + # Also ahead of the cost-table gate (the frac sweep is how the + # table gets built). Unlike the pin, the windows stay real: the + # forced fraction only replaces the argmax'd budget, and the + # confidence top-k spends it exactly as in production. + selected = self._gather_rows(num_gen_requests=num_gen_requests, + rows=rows) + if selected is None: + return None + self._note_snapshot_stats(selected, rows) + survival = compute_survival(self.apply_calibration(selected)) + trimmable = (self.cfg.resolved_max_verify_len + - self.cfg.min_verify_len) + budget = int(round(self._forced_budget_frac + * int(num_gen_requests) * trimmable)) + self.stats["forced_budget_steps"] = ( + self.stats.get("forced_budget_steps", 0) + 1) + lens = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=self.cfg).tolist() + else: + if self.cost_table is None or self.cost_table.is_flat: + self.stats["fallback_flat_cost"] += 1 + return None + selected = self._gather_rows(num_gen_requests=num_gen_requests, rows=rows) + if selected is None: + return None + + self._note_snapshot_stats(selected, rows) + probs = self.apply_calibration(selected) + survival = compute_survival(probs) + # Scored as the top-k redistribution that actually runs, on the + # same tier-aligned cost grid as the uniform argmax. + budget = compute_verify_token_budget( + survival=survival.numpy().astype(np.float64), + num_gen_requests=int(num_gen_requests), + cost_table=self.cost_table, + min_verify_len=self.cfg.min_verify_len, + max_verify_len=self.cfg.resolved_max_verify_len, + # Required: any non-tier total is rounded up to a captured + # bucket, pushing a budget chosen below a cost riser back over it. + allowed_lens=self.tiers, + ) + # Predicted step cost, recorded so runs can be reconciled against + # measured hostStepTimeMS; a systematic gap flags a wrong or + # mismatched cost table. + floor_tokens = int(num_gen_requests) * (self.cfg.min_verify_len + 1) + predicted = float(self.cost_table.step_times( + np.asarray([floor_tokens + int(budget)]), + int(num_gen_requests))[0]) + self.last_predicted_step_ms = predicted + self.stats["predicted_ms_sum"] = ( + self.stats.get("predicted_ms_sum", 0.0) + predicted) + self.stats["predicted_steps"] = ( + self.stats.get("predicted_steps", 0) + 1) + lens = schedule_verify_lens_topk(survival=survival, budget=budget, cfg=self.cfg).tolist() + # Local decision histograms, recorded before the bucket fit widens + # the realized windows (verify_len_hist is post-fit). + rung = self.cfg.min_verify_len + int(budget) // max(int(num_gen_requests), 1) + rhist = self.stats.setdefault("local_rung_hist", {}) + rhist[rung] = rhist.get(rung, 0) + 1 + lhist = self.stats.setdefault("local_len_hist", {}) + for v in lens: + lhist[int(v)] = lhist.get(int(v), 0) + 1 + + # The early returns above are rank-local (some timing-dependent), so no + # collective may run before this point; reduce once, unconditionally, + # at a point every rank reaches. reduce_across_ranks=False means the + # caller does that agreement itself. + if reduce_across_ranks: + all_rank_max = all_rank_max or self.all_rank_max + if all_rank_max is not None: + agreed_max = int(all_rank_max(int(max(lens)))) + lens = [min(int(v), agreed_max) for v in lens] + assert len(lens) == num_gen_requests, ( + f"internal: produced {len(lens)} verify lengths for {num_gen_requests} requests" + ) + return [int(v) for v in lens] diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 50e23d087ca7..0fc2168ad54f 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -132,8 +132,7 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata, return None if not isinstance(attn_metadata, TrtllmAttentionMetadata): return None - draft_block_offsets = getattr(attn_metadata, 'draft_kv_cache_block_offsets', - None) + draft_block_offsets = attn_metadata.draft_kv_cache_block_offsets if draft_block_offsets is None: return None @@ -518,6 +517,40 @@ class SpecMetadata: # Normally, it equals 1 + runtime_draft_len. But for PARD, it equals 2 * runtime_draft_len. runtime_tokens_per_gen_step: int = 1 + # Ragged (per-request) verification split, set only when + # confidence-scheduled verification runs in ragged mode: + # + # verify_lens [num_gen] positions verified per generation request + # qo_indptr [num_gen + 1] exclusive prefix sum; request r owns + # [qo_indptr[r], qo_indptr[r + 1]) + # + # Both stay None on every other path. Built by + # ``_torch/speculative/dspark_ragged.RaggedVerifyLayout``. + verify_lens: Optional[torch.Tensor] = None + qo_indptr: Optional[torch.Tensor] = None + # Host-side ``sum(verify_lens)`` (reading it off the tensor would sync). + # Equals the captured token bucket after ``RaggedVerifyLayout.fill_bucket``. + total_verify_tokens: Optional[int] = None + + # ``cap-accept``: per-request ceilings on the accepted count, in the same + # token units as ``verify_lens``, for a batch whose layout was NOT trimmed. + # Deliberately separate from ``verify_lens``: that field also implies the + # packed buffer layout, and under cap-accept the buffers are the usual + # rectangle -- reusing it would silently unpack a rectangle as if ragged. + accept_caps: Optional[torch.Tensor] = None + # [batch] positions the cap discarded, per request -- the acceptance the + # schedule cost each one. + accept_cap_trim: Optional[torch.Tensor] = None + + @property + def is_ragged_verify(self) -> bool: + """Whether this iteration uses per-request verify lengths. + + Implies the packed layout. ``cap-accept`` is deliberately NOT ragged by + this test -- it caps acceptance without packing anything. + """ + return self.verify_lens is not None + # Auto-detected per step from populated sampling params: # True if every request is greedy (no temp/top_k/top_p) and we can take # the argmax fast-path. False if any request needs sampling. @@ -659,8 +692,7 @@ def write_padding_onehot_draft_probs(self, padding_slot_ids, draft_len): overwritten here. No-op unless rejection is enabled and slots exist. Static shapes -> CUDA-graph safe. """ - if (not padding_slot_ids - or not getattr(self, "use_rejection_sampling", False) + if (not padding_slot_ids or not self.use_rejection_sampling or self.draft_probs is None): return onehot_vocab = (self.draft_probs_last_dim if self.draft_probs_last_dim @@ -716,7 +748,8 @@ def _scan_one_model_sampling( population. Does NOT allocate or fill GPU buffers, so it is safe to call before the CUDA graph key is built. """ - from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + from tensorrt_llm._torch.pyexecutor.llm_request import ( + LlmRequestState, get_request_tokens_per_gen_step) from tensorrt_llm._torch.pyexecutor.sampler.sampling_utils import \ GREEDY_TEMPERATURE_THRESHOLD from tensorrt_llm.sampling_params import SamplingParams @@ -784,8 +817,13 @@ def _normalize_request_sampling_params( tk_val = _first_or_none(sampling_config.top_k) tp_val = _first_or_none(sampling_config.top_p) - # Context requests have no draft tokens yet. - num_tokens = 1 + self.runtime_draft_len if request.state == LlmRequestState.GENERATION_IN_PROGRESS else 1 + # Generation requests contribute one row per verified position, + # per-request under ragged verification: the buffers are indexed + # by flat token position, so a fixed stride would apply request + # r's sampling params to request r+1's tokens. + num_tokens = get_request_tokens_per_gen_step( + request, 1 + self.runtime_draft_len + ) if request.state == LlmRequestState.GENERATION_IN_PROGRESS else 1 ( temp_val, @@ -997,6 +1035,59 @@ def populate_sampling_params_for_one_model( (tk for tk in request_top_ks if 0 < tk < _disable_topk), default=0) +def apply_accept_caps(num_accepted_tokens: torch.Tensor, num_contexts: int, + spec_metadata) -> None: + """``cap-accept``: hold each request to its scheduled window, in place. + + Must run inside acceptance, before the DSpark drafter consumes + ``num_accepted_tokens`` on device (hidden-state pick, rolling-KV + back-fill, ``_ctx_len`` advance); capping later leaves the draft state + permanently ahead of what was committed. In-place and sync-free, so safe + to record into a CUDA graph. + """ + caps = spec_metadata.accept_caps + if caps is None: + return + gen = num_accepted_tokens[num_contexts:] + if gen.numel() == 0: + return + caps = caps[:gen.shape[0]].to(gen.dtype) + capped = torch.minimum(gen, caps) + trim = spec_metadata.accept_cap_trim + if trim is not None: + # Zero the context slots too. The buffer is persistent and the batch + # composition moves every step, so a slot left unwritten would carry + # the previous step's loss and be attributed to whoever holds it now. + rows = num_contexts + gen.shape[0] + trim[:rows].zero_() + trim[num_contexts:rows] = (gen - capped).to(trim.dtype) + num_accepted_tokens[num_contexts:] = capped + + +def _padded_gen_draft_tokens(spec_metadata, num_gens: int, + runtime_draft_len: int) -> torch.Tensor: + """Unpack ragged-packed draft tokens into ``[num_gens, runtime_draft_len]``. + + The overlap path packs draft tokens by each request's own window, so the + buffer holds ``sum(verify_lens) - num_gens`` entries rather than a + ``num_gens x runtime_draft_len`` rectangle. Acceptance is naturally + rectangular, so scatter back out; positions past a request's window are + padding and get masked by ``count_accepted_ragged``. + """ + from .dspark_ragged import build_qo_indptr, scatter_ragged_to_padded + + # verify_lens counts tokens (bonus position + drafts); the draft buffer + # only holds the drafts. + draft_lens = spec_metadata.verify_lens - 1 + packed_len = spec_metadata.total_verify_tokens - num_gens + return scatter_ragged_to_padded( + spec_metadata.draft_tokens[:packed_len], + verify_lens=draft_lens, + qo_indptr=build_qo_indptr(draft_lens), + max_len=runtime_draft_len, + ).to(torch.int32) + + class SpecWorkerBase(nn.Module, ABC): """ Base class for speculative decoding workers. @@ -1373,16 +1464,37 @@ def _sample_and_accept_draft_tokens_base( accepted_tokens[:num_contexts, 0] = target_tokens[:num_contexts] # Generation requests: verify draft tokens against target tokens - gen_target_tokens = target_tokens[num_contexts:].reshape( - num_gens, runtime_draft_len + 1) - accepted_tokens[num_contexts:, :runtime_draft_len + - 1] = gen_target_tokens + if spec_metadata is not None and spec_metadata.is_ragged_verify: + from .dspark_ragged import (count_accepted_ragged, + scatter_ragged_to_padded) + verify_lens = spec_metadata.verify_lens + total = spec_metadata.total_verify_tokens + gen_target_tokens = scatter_ragged_to_padded( + target_tokens[num_contexts:num_contexts + total], + verify_lens=verify_lens, + qo_indptr=spec_metadata.qo_indptr, + max_len=runtime_draft_len + 1) + accepted_tokens[num_contexts:, :runtime_draft_len + + 1] = gen_target_tokens + # verify_lens counts tokens (bonus position included), so a + # request drafted verify_lens - 1 positions. Masking past that is + # what stops a stale padded slot from comparing equal and crediting + # an acceptance the target never made. + num_accepted_tokens[num_contexts:] += count_accepted_ragged( + draft_tokens=draft_tokens, + target_tokens=gen_target_tokens[:, :runtime_draft_len], + verify_lens=verify_lens - 1) + else: + gen_target_tokens = target_tokens[num_contexts:].reshape( + num_gens, runtime_draft_len + 1) + accepted_tokens[num_contexts:, :runtime_draft_len + + 1] = gen_target_tokens - # Compare draft tokens with target tokens using cumulative product - # Counts consecutive matches from the start - num_accepted_tokens[num_contexts:] += torch.cumprod( - (draft_tokens == gen_target_tokens[:, :runtime_draft_len]).int(), - dim=-1).sum(1) + # Compare draft tokens with target tokens using cumulative product + # Counts consecutive matches from the start + num_accepted_tokens[num_contexts:] += torch.cumprod(( + draft_tokens == gen_target_tokens[:, :runtime_draft_len]).int(), + dim=-1).sum(1) # Apply force override if set num_accepted_tokens = self._apply_force_accepted_tokens( @@ -1391,6 +1503,16 @@ def _sample_and_accept_draft_tokens_base( runtime_draft_len, spec_metadata=spec_metadata) + if spec_metadata is not None and spec_metadata.is_ragged_verify: + # The synthetic-acceptance override caps at the batch-wide + # runtime_draft_len; under ragged a request may have verified + # fewer positions, and accepting past its window would commit + # tokens the target never produced. + num_accepted_tokens[num_contexts:] = torch.minimum( + num_accepted_tokens[num_contexts:], + spec_metadata.verify_lens.to(num_accepted_tokens.dtype)) + apply_accept_caps(num_accepted_tokens, num_contexts, spec_metadata) + return accepted_tokens, num_accepted_tokens def _accept_draft_tokens(self, logits, draft_tokens, num_contexts, @@ -1444,7 +1566,7 @@ def _draft_logits_are_sharded(self, logits, spec_metadata): sharded and gather identical copies (overflowing the d2t table). """ mapping = self.mapping - if mapping is None or getattr(mapping, "tp_size", 1) <= 1: + if mapping is None or mapping.tp_size <= 1: return False # Under attention DP every rank is data-parallel -- it owns a distinct # set of requests -- so the draft logits must NOT be cross-rank gathered, @@ -1455,7 +1577,7 @@ def _draft_logits_are_sharded(self, logits, spec_metadata): # desyncing them into a hang / DeepEP launch failure. Only plain TP # (tp>1 without ADP), where all ranks share the same tokens sharded over # the vocab dim, needs the gather. - if getattr(mapping, "enable_attention_dp", False): + if mapping.enable_attention_dp: return False draft_full_vocab = (getattr(spec_metadata, "draft_vocab_size", 0) or getattr(spec_metadata, "vocab_size", 0) or 0) @@ -1552,6 +1674,12 @@ def _reshape_draft_tokens_for_accept(self, spec_metadata, num_gens, device): return torch.zeros((num_gens, runtime_draft_len), dtype=torch.int, device=device) + if spec_metadata.is_ragged_verify: + # Packed by per-request windows, not a rectangle; a plain reshape + # can happen to divide and silently attribute one request's drafts + # to another. + return _padded_gen_draft_tokens(spec_metadata, num_gens, + runtime_draft_len) return spec_metadata.draft_tokens.reshape(num_gens, runtime_draft_len) def _reshape_logits_for_accept(self, logits, num_contexts, num_gens, @@ -1608,9 +1736,22 @@ def _rejection_buffers_valid(self, draft_tokens, draft_len, stored_vocab, return False if draft_tokens.dim() != 2 or draft_tokens.shape[0] != num_gens: return False - # logits must cover context rows (1 each) + gen rows (draft_len + 1 each). + # logits must cover context rows (1 each) + however many gen rows the + # target was actually given: under ragged verification that is + # `total_verify_tokens`, not `num_gens * (draft_len + 1)`. logits_rows = logits.shape[0] if logits.dim() > 1 else 1 - if logits_rows < num_contexts + num_gens * (draft_len + 1): + if spec_metadata.is_ragged_verify: + # Fail closed on a half-built ragged layout: the branch below + # dereferences all three of these. + total_verify_tokens = spec_metadata.total_verify_tokens + if (total_verify_tokens is None + or spec_metadata.verify_lens is None + or spec_metadata.qo_indptr is None): + return False + required_gen_rows = int(total_verify_tokens) + else: + required_gen_rows = num_gens * (draft_len + 1) + if logits_rows < num_contexts + required_gen_rows: return False # Slot ids for the gen subset must exist (range safety is guaranteed by # construction). @@ -1674,7 +1815,12 @@ def _sample_and_accept_draft_tokens_rejection( # === Generation subset: rejection sampling on the gen slice === if num_gens > 0: - num_gen_logits = num_gens * (runtime_draft_len + 1) + is_ragged = spec_metadata.is_ragged_verify + # The per-token sampling-parameter buffers are packed by the same + # per-request windows as the logits (see + # _scan_one_model_sampling), so one flat total covers both. + num_gen_logits = (spec_metadata.total_verify_tokens if is_ragged + else num_gens * (runtime_draft_len + 1)) gen_logits = logits[num_contexts:num_contexts + num_gen_logits] gen_start = num_contexts gen_end = num_contexts + num_gen_logits @@ -1690,9 +1836,19 @@ def _sample_and_accept_draft_tokens_rejection( target_probs_flat = compute_probs_from_logits( gen_logits, temperatures, top_ks, top_ps) - target_probs = target_probs_flat.reshape(num_gens, - runtime_draft_len + 1, - vocab_size) + if is_ragged: + from .dspark_ragged import (fill_padded_rows_onehot, + scatter_ragged_to_padded) + target_probs = scatter_ragged_to_padded( + target_probs_flat, + verify_lens=spec_metadata.verify_lens, + qo_indptr=spec_metadata.qo_indptr, + max_len=runtime_draft_len + 1) + fill_padded_rows_onehot(target_probs, + verify_lens=spec_metadata.verify_lens) + else: + target_probs = target_probs_flat.reshape( + num_gens, runtime_draft_len + 1, vocab_size) draft_vocab_size = draft_probs.shape[-1] assert draft_probs.shape[0] == num_gens, ( @@ -1776,6 +1932,16 @@ def _sample_and_accept_draft_tokens_rejection( num_contexts, runtime_draft_len, spec_metadata=spec_metadata) + + if num_gens > 0 and spec_metadata.is_ragged_verify: + # The kernel walks the padded rectangle, so a request whose whole + # window was accepted keeps going into padding. verify_lens is the + # hard ceiling: accepting past it commits tokens the target never + # scored. + num_accepted_tokens[num_contexts:] = torch.minimum( + num_accepted_tokens[num_contexts:], + spec_metadata.verify_lens.to(num_accepted_tokens.dtype)) + apply_accept_caps(num_accepted_tokens, num_contexts, spec_metadata) return accepted_tokens, num_accepted_tokens def _update_advance_draft_sampling_seed(self, device): @@ -1859,7 +2025,7 @@ def greedy_sample_draft_with_tp_gather(self, ``[:batch]``, which belong to group rank 0. """ if (mapping_lm_head_tp is not None - and getattr(mapping_lm_head_tp, "tp_size", 1) > 1): + and mapping_lm_head_tp.tp_size > 1): from ..distributed.ops import allgather combined = self._get_local_max_and_combined(logits, mapping_lm_head_tp) @@ -1871,8 +2037,7 @@ def greedy_sample_draft_with_tp_gather(self, return self._get_draft_tokens_from_gathered(own_segment) mapping = self.mapping sharded = self._draft_logits_are_sharded(logits, spec_metadata) - if (sharded and mapping is not None - and getattr(mapping, "tp_size", 1) > 1 + if (sharded and mapping is not None and mapping.tp_size > 1 and not mapping.enable_attention_dp): from ..distributed.ops import allgather combined = self._get_local_max_and_combined(logits) @@ -1917,7 +2082,7 @@ def advanced_sample_draft_block(self, gen_logits: torch.Tensor, self._update_advance_draft_sampling_seed(gen_logits.device) flat_logits = gen_logits.reshape(num_gens * K, vocab) - if getattr(spec_metadata, "use_rejection_sampling", False): + if spec_metadata.use_rejection_sampling: flat_tokens, flat_probs = ( sampling_batch_spec_dec_one_model_for_rejection( flat_logits, @@ -1966,8 +2131,7 @@ def write_context_onehot_draft_probs(self, next iter reads this scalar). No-op unless rejection is enabled and the slot-indexed buffers exist. Static shapes -> CUDA-graph safe. """ - if (num_contexts <= 0 - or not getattr(spec_metadata, "use_rejection_sampling", False) + if (num_contexts <= 0 or not spec_metadata.use_rejection_sampling or spec_metadata.draft_probs is None): return ctx_slot_ids = spec_metadata.batch_slot_ids[:num_contexts] @@ -2136,8 +2300,7 @@ def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager): return # Check if draft KV cache block offsets are allocated - draft_block_offsets = getattr(attn_metadata, - 'draft_kv_cache_block_offsets', None) + draft_block_offsets = attn_metadata.draft_kv_cache_block_offsets if draft_block_offsets is None: # Draft KV cache block offsets not allocated, skip switching yield diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 3d841aaa4a38..95996ca16f62 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -47,6 +47,13 @@ class SampleStateTensorsSpec(SampleStateTensors): new_tokens_lens: torch.Tensor next_draft_tokens: torch.Tensor + #: cap-accept only; None on every other path. + cap_trim_lens: Optional[torch.Tensor] = None + #: ragged verification only; None on every other path. The TOKEN windows + #: the step actually verified (bonus included), slot-indexed. Under + #: device-window selection this D2H is the host's only source of the true + #: windows; zero marks "no window this step". + verify_lens: Optional[torch.Tensor] = None @dataclass(kw_only=True) @@ -55,6 +62,26 @@ class SampleStateSpec(SampleState): device: SampleStateTensorsSpec host: SampleStateTensorsSpec + #: This step's per-request verify windows, keyed by request id. Snapshotted + #: because ``update_requests`` for step N-1 runs *after* the scheduler has + #: already stamped step N's windows onto the same LlmRequest objects, so + #: reading ``request.py_verify_len`` live rewinds by the wrong amount under + #: the overlap scheduler. None on every uniform path. + verify_lens_snapshot: Optional[dict] = None + #: Same idea, for ``cap-accept``: the window the scheduler chose, which + #: bounds *commitment* only -- the target still scored the full block, so + #: the KV rewind must use the full block and not this. Kept in a separate + #: field from ``verify_lens_snapshot`` for exactly that reason: the two + #: never coexist, and conflating them would rewind by the window after + #: verifying the block, silently leaking KV. + verify_caps_snapshot: Optional[dict] = None + #: The draft pass that produced the block THIS state's tokens verified; + #: the STS recorder joins its acceptance label to a logits snapshot by + #: this key. Snapshotted at sampling time for the same reason + #: ``verify_lens_snapshot`` is: by the time ``update_requests`` runs, the + #: live counter already belongs to a later pass. None unless STS + #: collection is on. + sts_target_seq: Optional[int] = None class SpecSamplerBase(Sampler[SampleStateSpec], AsyncWorkerMixin): @@ -85,6 +112,15 @@ class Store: next_new_tokens: torch.Tensor next_draft_tokens: torch.Tensor new_tokens_lens: torch.Tensor + #: cap-accept: per-request positions the verify window discarded. + #: Slot-indexed like new_tokens_lens, and rewritten for every slot in + #: the batch each step -- a slot left alone would hand the previous + #: occupant's loss to whoever holds it now. + cap_trim_lens: torch.Tensor + #: ragged verification: the TOKEN windows the step actually verified, + #: slot-indexed; zeroed for every batch slot on non-ragged steps (the + #: same stale-slot discipline as cap_trim_lens). + verify_lens: torch.Tensor def __init__(self, args: TorchSampler.Args, *, draft_len: int): """ @@ -111,6 +147,8 @@ def __init__(self, args: TorchSampler.Args, *, draft_len: int): next_new_tokens=int_tensor((max_tokens, seq_slots, self.max_beam_width)), next_draft_tokens=int_tensor((seq_slots, draft_tokens_size)), new_tokens_lens=int_tensor((seq_slots,)), + cap_trim_lens=int_tensor((seq_slots,)), + verify_lens=int_tensor((seq_slots,)), ) def _get_max_tokens(self, args: TorchSampler.Args, draft_len: int) -> int: @@ -175,6 +213,56 @@ def _request_common_handling( request.py_draft_tokens = next_draft_tokens[request.py_seq_slot][:runtime_draft_len] request.py_decoding_iter += 1 + #: Optional observability sink, attached by the DSpark worker. Left None + #: everywhere else so this costs one attribute read per request. + acceptance_stats = None + #: STS calibration collection; attached alongside `acceptance_stats` and + #: None everywhere else. `sts_row_for` resolves a py_request_id to its row + #: in the worker's confidence buffer through the worker's OWN allocator -- + #: `py_seq_slot` comes from a different allocator whose numbering only + #: coincides until the first request completes. `sts_seq_provider` names + #: the draft pass the current step's verification targets. + sts_recorder = None + sts_row_for = None + sts_seq_provider = None + + @staticmethod + def _verified_len(request: LlmRequest, runtime_draft_len: int, + verify_lens_snapshot: Optional[dict] = None, + ridden_verify_lens: Optional[list] = None) -> int: + """How many draft positions THIS request was given to verify. + + Uniform scheduling gives every request the batch-wide + ``runtime_draft_len``; ragged (per-request) verification does not, and + the KV rewind below is computed per request: + + rewind = verified_positions - accepted_positions + + Using the batch-wide value under ragged verification rewinds the wrong + amount for every request that got a shorter window -- silently dropping + or keeping KV entries, i.e. wrong output with no error anywhere. + + ``py_verify_len`` is set only by the ragged path, so this is exactly + today's behavior for every other speculation mode. + """ + # First preference: the windows the device layout actually verified, + # ridden back on this step's sampler D2H (slot-indexed TOKEN windows; + # zero = no window). Under device-window selection the host attribute + # holds only a shape split, so this is the only correct source; under + # host windows the two agree. + if ridden_verify_lens is not None: + token_window = ridden_verify_lens[request.py_seq_slot] + if token_window > 0: + return int(token_window) - 1 + # Second: the snapshot taken when this step's tokens were sampled. The + # live attribute already belongs to the NEXT step by the time the + # overlap scheduler rewinds this one. + if verify_lens_snapshot is not None: + per_request = verify_lens_snapshot.get(request.py_request_id) + else: + per_request = getattr(request, "py_verify_len", None) + return runtime_draft_len if per_request is None else int(per_request) + def update_requests( self, state: SampleStateSpec, @@ -193,14 +281,25 @@ def update_requests( state.sampler_event.synchronize() new_tokens = state.host.new_tokens.tolist() new_tokens_lens_list = state.host.new_tokens_lens.tolist() + cap_trim_list = (state.host.cap_trim_lens.tolist() + if state.host.cap_trim_lens is not None else None) + ridden_verify_lens = (state.host.verify_lens.tolist() + if state.host.verify_lens is not None else None) next_draft_tokens_list = state.host.next_draft_tokens.tolist() beam_idx = DEFAULT_BEAM_IDX runtime_draft_len = getattr(state, "runtime_draft_len", self.draft_len) + caps = state.verify_caps_snapshot for req in state.requests: if req.state == LlmRequestState.GENERATION_COMPLETE: continue num_new_tokens = new_tokens_lens_list[req.py_seq_slot] + # cap-accept: `num_new_tokens` has ALREADY been capped, on device, + # inside acceptance (`interface.apply_accept_caps`) -- it has to be, + # because the drafter reads the same count later in that forward to + # advance its own state. The cap is read here only to report the + # window the request was actually given. + cap = None if caps is None else caps.get(req.py_request_id) for i in range(num_new_tokens): new_token = add_token(req, new_tokens, beam_idx=beam_idx, step=i) if TorchSampler._handle_stop_criteria( @@ -208,9 +307,47 @@ def update_requests( ): break req.py_num_accepted_draft_tokens = num_new_tokens - 1 - req.py_rewind_len = runtime_draft_len - req.py_num_accepted_draft_tokens + verified_len = self._verified_len(req, runtime_draft_len, + state.verify_lens_snapshot, + ridden_verify_lens) + req.py_rewind_len = verified_len - req.py_num_accepted_draft_tokens + # Acceptance against the window the request was actually given. + # Recorded here because this is the only place both numbers are + # host-side and belong to the same step -- the live py_verify_len + # has already moved on by now, which is why the snapshot exists. + # `acceptance_stats` is attached by the DSpark worker and is None + # for every other path. + # + # Under cap-accept the window is the cap, NOT `verified_len`: the + # target verified the whole block (which is what makes the rewind + # above correct), but the scheduler only granted `cap`. Reporting + # the block here would count every request as untrimmed and hide + # the mode's entire measurement. + if self.acceptance_stats is not None: + self.acceptance_stats.record_acceptance( + accepted=req.py_num_accepted_draft_tokens, + window=verified_len if cap is None else int(cap), + cap_trim=(0 if cap_trim_list is None else + cap_trim_list[req.py_seq_slot])) + # STS calibration pairs this request's drafted-block logits with + # how much of that block was accepted. The logits come from the + # recorder's snapshot ring, selected by the draft pass snapshotted + # into this state at sampling time and checked against the row's + # own stamp -- reading the live buffer here paired the label with + # whatever a LATER pass had written over it. + if self.sts_recorder is not None: + self.sts_recorder.record( + row=(self.sts_row_for(req.py_request_id) + if self.sts_row_for is not None else None), + accepted=req.py_num_accepted_draft_tokens, + target_seq=state.sts_target_seq) self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) + if self.sts_recorder is not None: + # Periodic flush, so a run killed mid-collection still leaves usable + # shards rather than everything living in a list that dies with it. + self.sts_recorder.end_step() + def sample_async( self, scheduled_requests: ScheduledRequests, @@ -279,6 +416,29 @@ def sample_async( self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens) self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens) self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens) + # cap-accept: absent on every other path, and written as ZEROS then + # rather than skipped. The buffer is persistent and slot-indexed, so a + # slot left untouched would keep whatever the previous occupant lost + # and hand it to the current one. + o_cap_trim = outputs.get("cap_trim_lens") + if o_cap_trim is None: + o_cap_trim = torch.zeros_like(o_new_tokens_lens) + else: + o_cap_trim = o_cap_trim[num_skip:num_skip + num_sampling_requests] + self.store.cap_trim_lens.index_copy_( + 0, slots, o_cap_trim.to(self.store.cap_trim_lens.dtype)) + # Ragged verification: the true token windows, same zero-then-write + # slot discipline as cap_trim_lens. Under device-window selection the + # host py_verify_len holds only a shape split, so the rewind + # arithmetic in update_requests must read THIS. + o_verify_lens = outputs.get("verify_lens") + if o_verify_lens is None: + o_verify_lens = torch.zeros_like(o_new_tokens_lens) + else: + o_verify_lens = o_verify_lens[num_skip:num_skip + + num_sampling_requests] + self.store.verify_lens.index_copy_( + 0, slots, o_verify_lens.to(self.store.verify_lens.dtype)) # Create sample state with async D2H copy device_tensors = SampleStateTensorsSpec( @@ -291,6 +451,10 @@ def sample_async( new_tokens=self._copy_to_host(self.store.new_tokens), new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens), next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens), + # [max_num_sequences] int32 -- kilobytes next to new_tokens above, + # and it rides the same event, so no extra synchronization. + cap_trim_lens=self._copy_to_host(self.store.cap_trim_lens), + verify_lens=self._copy_to_host(self.store.verify_lens), ) sampler_event = self._record_sampler_event() @@ -299,10 +463,29 @@ def sample_async( for request in finished_context_requests: request.py_draft_tokens = [1] * self.draft_len + # Cheap (a few ints) and only populated when the ragged path is live. + verify_lens_snapshot = None + verify_caps_snapshot = None + for request in sampling_requests: + window = getattr(request, "py_verify_len", None) + if window is not None: + if verify_lens_snapshot is None: + verify_lens_snapshot = {} + verify_lens_snapshot[request.py_request_id] = int(window) + cap = getattr(request, "py_verify_cap", None) + if cap is not None: + if verify_caps_snapshot is None: + verify_caps_snapshot = {} + verify_caps_snapshot[request.py_request_id] = int(cap) + return SampleStateSpec( requests=sampling_requests, device=device_tensors, host=host_tensors, sampler_event=sampler_event, runtime_draft_len=runtime_draft_len, + verify_lens_snapshot=verify_lens_snapshot, + verify_caps_snapshot=verify_caps_snapshot, + sts_target_seq=(self.sts_seq_provider() + if self.sts_seq_provider is not None else None), ) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index b077747b6c03..adcc76a2945a 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -51,6 +51,16 @@ LM_EVAL_DEFAULT_IMAGE_PLACEHOLDER = "" +def _is_metric_key(key: str) -> bool: + """Whether a ``results[task]`` key is a metric rather than bookkeeping. + + lm-eval keys every metric as ``","`` -- ``"exact_match, + strict-match"``, ``"acc,none"``. Entries without a comma are bookkeeping: + ``"alias"`` is the task name, ``"samples"`` is the example count. + """ + return "," in key + + class LmEvalWrapper(TemplateLM): def __init__(self, @@ -635,10 +645,11 @@ def evaluate(self, system_instruction=self.system_prompt, log_samples=self.log_samples) - # Normalize scores to range 0~100 + # Normalize real metric scores to range 0~100 (see _is_metric_key). scores = results["results"][self.task_name] for metric in scores.keys(): - if isinstance(scores[metric], (float, int)): + if _is_metric_key(metric) and isinstance(scores[metric], + (float, int)): scores[metric] *= 100 logger.info( f"lm-eval {self.task_name} results (scores normalized to range 0~100):\n{lm_eval.utils.make_table(results)}" @@ -654,8 +665,18 @@ def evaluate(self, f"lm-eval {self.task_name} {scores_filter} accuracy: {result_acc:.2f}" ) else: - result_acc = np.mean( - [acc for m, acc in scores.items() if "_stderr" not in m]) + # Average real metrics only; bookkeeping keys ("alias", + # "samples") are excluded structurally (see _is_metric_key). + numeric = [ + acc for m, acc in scores.items() + if _is_metric_key(m) and "_stderr" not in m + and isinstance(acc, (int, float)) and not isinstance(acc, bool) + ] + if not numeric: + raise ValueError( + f"lm-eval returned no numeric metrics for " + f"{self.task_name}; got {sorted(scores)}") + result_acc = np.mean(numeric) logger.info( f"lm-eval {self.task_name} average accuracy: {result_acc:.2f}") return result_acc diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 5e1413a4f0da..de0f774c8ded 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -872,6 +872,32 @@ def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: f"{action}() failed on {len(errors)} rank(s):\n" + "\n".join(errors)) + def set_dspark_verify_len_pin(self, + verify_len: Optional[int]) -> Optional[int]: + """Pin (or unpin) the DSpark verify length on this worker's executor. + + Safe under TP/PP without the control-communicator dance ``sleep`` needs: + the pin is a scalar the decode loop reads on its next step, and it is + broadcast to the other ranks by the step's own allgather rather than by + this call. So rank 0 alone receiving the RPC is sufficient and the loop + never has to be paused. + """ + if self.engine is None: + raise RuntimeError("engine is not initialized") + return self.engine.set_dspark_verify_len_pin(verify_len) + + def set_dspark_budget_frac(self, + frac: Optional[float]) -> Optional[float]: + """Set (or clear) the DSpark verify-budget fraction on this executor. + + Same delivery contract as :meth:`set_dspark_verify_len_pin`: a scalar + queued on rank 0 and handed to every rank by the decode loop's own + per-step allgather. + """ + if self.engine is None: + raise RuntimeError("engine is not initialized") + return self.engine.set_dspark_budget_frac(frac) + def sleep(self, sleep_tags: list[str]) -> None: """Release GPU virtual memory for the specified memory type tags. @@ -1000,6 +1026,7 @@ def _stats_serializer(stats) -> str: prev_device_step_time_ms = stats[5] if len(stats) > 5 else None scheduler_mode = stats[6] if len(stats) > 6 else None gpu_forward_time_ms = stats[7] if len(stats) > 7 else None + num_generation_tokens = stats[8] if len(stats) > 8 else None stats_dict = json.loads(iteration_stats.to_json_str()) # Always tag the row so Dynamo's adapter can read @@ -1042,6 +1069,11 @@ def _stats_serializer(stats) -> str: # CPU wall. Set per-record so consumers do not need server config. if scheduler_mode is not None: stats_dict["schedulerMode"] = scheduler_mode + # Executed (padded) generation-token total -- the M axis of the DSpark + # SPS cost table. Rank 0's value, broadcast per row by the ADP fanout; + # group-uniform on graph steps by construction. + if num_generation_tokens is not None: + stats_dict["numGenerationTokens"] = int(num_generation_tokens) # Convert back to JSON string return json.dumps(stats_dict) diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index d1baa0792785..f9cf259ea43a 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -61,7 +61,11 @@ # sleep() and wakeup() coordinate across all ranks via a dedicated control # communicator started in PyExecutor.start_worker(); rank-0 is the entry # point, so routing through the rank-0 RPC shim is correct. -_MULTI_RANK_ALLOWED_METHODS: frozenset[str] = frozenset({"sleep", "wakeup"}) +# set_dspark_verify_len_pin() needs no control communicator at all: it queues +# a scalar that the decode loop's own per-step allgather hands to every rank, +# so rank-0 receiving it is sufficient and the loop is never paused. +_MULTI_RANK_ALLOWED_METHODS: frozenset[str] = frozenset( + {"sleep", "wakeup", "set_dspark_verify_len_pin", "set_dspark_budget_frac"}) def _check_collective_rpc_guard( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 68d9fdbeb207..45ecf39443a0 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2759,13 +2759,13 @@ class DSparkDecodingConfig(DecodingBaseConfig): layers as cross-attention context and drafts a whole block in a single backbone forward, but it additionally refines the per-position draft logits with a lightweight sequential head (a low-rank Markov head, optionally an RNN - head) and predicts an acceptance-confidence per position to truncate the - proposed prefix. + head) and predicts a per-position acceptance confidence used to budget how + much verification the batch is worth. Key features: - Target-dependent: captures hidden states from ``target_layer_ids``. - Semi-parallel: one block backbone forward + cheap sequential head refine. - - Confidence head: truncates the proposed draft length (NOT the accept rule; + - Confidence head: sizes the *verification* budget (NOT the accept rule; acceptance stays standard target verification, preserving greedy parity). Reference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec). @@ -2802,16 +2802,146 @@ class DSparkDecodingConfig(DecodingBaseConfig): "read from the draft model config (dspark_markov_head_type), " "defaulting to \"vanilla\".") - # NOTE: confidence-based dynamic drafting (the draft model's confidence head - # that truncates the proposed block) is NOT enabled in this PR. The user-facing - # ``enable_confidence_head`` / ``confidence_threshold`` knobs are intentionally - # omitted and will be added when the feature is actually wired into the - # speculative scheduling/verification path. The confidence head module and its - # internal plumbing remain as scaffolding (see DSparkConfidenceHead / - # dspark_propose). + enable_confidence_scheduling: bool = Field( + default=False, + description= + "Score each drafted position with the draft's confidence head and let " + "the verification scheduler size the batch's verify budget from those " + "scores. The block is always drafted in full; only the number of tokens " + "sent to the target for verification changes, so acceptance (and hence " + "output distribution) is unaffected. Requires a checkpoint with trained " + "confidence-head weights. Gains only appear at high concurrency -- at " + "small batch sizes the step cost is nearly flat in the verified token " + "count, so trimming saves nothing and costs acceptance length.", + status="prototype") + + confidence_sts_path: Optional[str] = Field( + default=None, + description= + "Path to a JSON file with per-position sequential-temperature-scaling " + "(STS) values used to calibrate the confidence head. The scheduler " + "consumes a cumulative product of per-position probabilities, so a " + "per-position calibration error compounds geometrically along the " + "block. If None, the raw sigmoid is used (identity calibration).", + status="prototype") + + confidence_sps_table_path: Optional[str] = Field( + default=None, + description= + "Path to a JSON file with the profiled decode step cost as a function of " + "the total verified token count. Without it the planner has a flat cost " + "model, under which every extra verified token looks free and the budget " + "degenerates to verify-all (no scheduling gain).", + status="prototype") + + confidence_verify_len_tiers: Optional[List[PositiveInt]] = Field( + default=None, + description= + "Draft lengths to capture CUDA graphs for. The scheduler may only pick " + "from this ladder: a length with no captured graph does not raise, it " + "silently drops the step out of graph replay into eager execution, which " + "costs far more than the trimmed tokens save. Each extra tier is another " + "captured graph, and captured graphs consume memory that would otherwise " + "be KV cache, so keep the ladder short. If None, defaults to " + "[1, ceil(max_draft_len/2), max_draft_len].", + status="prototype") + + enable_ragged_verify: bool = Field( + default=False, + description= + "Give each request its own verify length instead of one length for the " + "whole batch. The batch's verify budget is then allocated by global " + "survival rank, so a request whose draft is still likely alive at depth " + "5 can take positions from one whose draft already died at depth 1. " + "Requires enable_confidence_scheduling. The packed batch is padded to a " + "captured (batch size, token count) pair, so this captures a second " + "graph axis and costs more graph memory than uniform scheduling.", + status="prototype") decoding_type: Literal["DSpark"] = Field(default="DSpark") + @model_validator(mode="after") + def validate_ragged_verify(self): + if self.enable_ragged_verify and not self.enable_confidence_scheduling: + raise ValueError( + "enable_ragged_verify requires enable_confidence_scheduling=True: " + "per-request verify lengths are chosen from confidence-head " + "survival, which is only computed when scheduling is on") + if self.enable_confidence_scheduling and not self.enable_ragged_verify: + # The two flags describe one feature. Scheduling without ragged used + # to mean a uniform tier ladder -- one window chosen for the whole + # batch -- which is a middle ground that does not earn its keep: + # it cannot give a confident request a longer window than a + # collapsing one, which is the entire point of scheduling against + # per-request confidence, and it degenerates to the full block + # whenever acceptance is high (every workload measured on + # DeepSeek-V4-Pro-DSpark). SGLang has no equivalent either: its + # scheduler always produces per-request windows via top-k, and + # uniform appears only as the verify-all degenerate case. + # + # So there are two states, not three: schedule per request, or + # verify the full block. Declining the feature is + # enable_confidence_scheduling=False. + raise ValueError( + "enable_confidence_scheduling requires enable_ragged_verify=True. " + "Scheduling a single window for the whole batch cannot act on " + "per-request confidence, and collapses to verifying the full " + "block whenever acceptance is high. To turn the feature off, " + "set enable_confidence_scheduling=False -- that verifies the " + "full drafted block, which is the only fallback.") + if (self.enable_ragged_verify and not self.confidence_sps_table_path + and not self.confidence_sts_path): + # Without a profiled cost table the planner's budget degenerates to + # verify-all -- correctly, since a flat model makes every extra + # token look free -- so the ragged path silently never runs. + raise ValueError( + "enable_ragged_verify=True requires a profiled cost table: pass " + "confidence_sps_table_path (produced by " + "`python tests/microbenchmarks/dspark_sps_profiler.py`) " + "or confidence_sts_path. Without one the planner cannot compare " + "the cost of extra verify tokens against their expected " + "acceptance, so it declines to trim and every request receives " + "the full window -- the ragged path runs but changes nothing. " + "Collecting a table does NOT require this flag: the profiler " + "pins static mode and sweeps the verify length by rebuilding " + "the engine, so run it with ragged verification off.") + return self + + @model_validator(mode="after") + def validate_confidence_scheduling(self): + if not self.enable_confidence_scheduling: + if (self.confidence_sts_path or self.confidence_sps_table_path + or self.confidence_verify_len_tiers): + raise ValueError( + "confidence_sts_path / confidence_sps_table_path / " + "confidence_verify_len_tiers require " + "enable_confidence_scheduling=True") + return self + + if self.max_draft_len is not None: + tiers = self.confidence_verify_len_tiers or [ + 1, max(1, (self.max_draft_len + 1) // 2), self.max_draft_len + ] + tiers = sorted({int(t) for t in tiers}) + if tiers[-1] > self.max_draft_len: + raise ValueError( + f"confidence_verify_len_tiers {tiers} exceeds max_draft_len " + f"({self.max_draft_len}); a request cannot verify more tokens " + f"than the draft proposes") + # The full block must stay reachable: it is the fallback used whenever + # the confidence snapshot is stale or the cost model is unprofiled. + if tiers[-1] != self.max_draft_len: + tiers.append(self.max_draft_len) + self.confidence_verify_len_tiers = tiers + return self + + @property + def verify_len_tiers(self) -> List[int]: + """Captured draft-length ladder; ``[max_draft_len]`` when scheduling is off.""" + if not self.enable_confidence_scheduling: + return [self.max_draft_len] + return list(self.confidence_verify_len_tiers or [self.max_draft_len]) + @model_validator(mode="after") def set_max_total_draft_tokens(self): self.max_total_draft_tokens = self.max_draft_len @@ -5392,6 +5522,73 @@ def validate_encode_only_torch_compile_config(self) -> 'TorchLlmArgs': "graphs or disable enable_piecewise_cuda_graph.") return self + def _validate_dspark_ragged_verify_environment(self) -> None: + """Reject the configurations ragged verification cannot honour. + + Every check here guards a path that would otherwise degrade *silently*: + the run comes up, produces plausible text, and either never takes the + ragged path at all or takes it with a kernel that mis-attributes rows. + A hard error at construction is the only signal that survives. + """ + cuda_graph_config = self.cuda_graph_config + # `fit_ragged_verify_lens` plans against the padded row count and + # reserves tokens for the CUDA-graph padding rows, so it presumes those + # rows exist. With padding off, `_get_padded_batch` adds none: the batch + # then carries a token total that keys into a graph nobody captured and + # drops to eager on every ragged step -- which costs far more than the + # trimmed tokens save -- while real requests were still shrunk to leave + # room for padding that never arrived. + if cuda_graph_config is None: + raise ValueError( + "speculative_config.enable_ragged_verify=True requires " + "cuda_graph_config to be set with enable_padding=True. Ragged " + "verification fits each step's token total onto a captured " + "CUDA-graph bucket; with no graphs there is no bucket to fit " + "and every step would silently run eager.") + if not cuda_graph_config.enable_padding: + raise ValueError( + "speculative_config.enable_ragged_verify=True requires " + "cuda_graph_config.enable_padding=True (it defaults to False). " + "The ragged token budget is planned against the padded batch " + "size and reserves tokens for the padding rows; without padding " + "those rows never appear, so the step's token total misses its " + "captured bucket and falls out of graph replay.") + + sparse_attention_config = self.sparse_attention_config + if sparse_attention_config is not None: + # Both recover "which request is this row" from a fixed next_n -- + # GVR indexes its previous-step top-k hint by request and strides it + # by next_n, and the CuTe-DSL kernel takes next_n as a compile-time + # constant and JIT cache key. Neither can express per-request + # windows. The C++ launcher already refuses the combination + # (indexerTopK.cu, "ragged rowKvLens is incompatible with the GVR + # preIdx hint"), but that fires deep in a decode step; say so here. + incompatible = [ + name for name in ("enable_heuristic_topk", "use_cute_dsl_topk") + if getattr(sparse_attention_config, name, False) + ] + if incompatible: + raise ValueError( + "speculative_config.enable_ragged_verify=True is " + "incompatible with sparse_attention_config." + f"{' / '.join(incompatible)}=True: those top-k paths " + "reconstruct each row's request and window position from a " + "single batch-wide next_n, which a ragged batch does not " + "have. Disable them, or disable ragged verification.") + # Not an error: the DSL paged-MQA-logits path is structurally + # uniform too, but the expanded (one row per query token) path it + # falls back to is numerically equivalent, and it is the default on + # DeepSeek-V4 anyway. Still worth saying out loud, since the knob + # will read as honoured otherwise. + if getattr(sparse_attention_config, + "use_cute_dsl_paged_mqa_logits", False): + logger.warning( + "sparse_attention_config.use_cute_dsl_paged_mqa_logits=True " + "is ignored on ragged-verification steps: the DSL kernel " + "broadcasts one KV length per request across a fixed next_n. " + "Those steps use the expanded DeepGEMM path instead, which " + "is numerically equivalent.") + @model_validator(mode="after") def validate_speculative_config(self): if self.speculative_config: @@ -5650,6 +5847,8 @@ def _dspark_get(key, top_level_key): "DSpark block_size must equal max_draft_len; got " f"block_size={spec_cfg.block_size} and " f"max_draft_len={spec_cfg.max_draft_len}") + if spec_cfg.enable_ragged_verify: + self._validate_dspark_ragged_verify_environment() if isinstance(self.speculative_config, SADecodingConfig): pool_size = self.speculative_config.global_pool_size diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 303d79e9ee13..9be5fb87e0da 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -866,7 +866,18 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo indices.extend([BAD_PAGE_INDEX] * (new_num_blocks - old_num_blocks)) else: if len(indices) < new_num_blocks: - raise ValueError("User-provided base page indices is too short") + # The buffer is a slice of host_kv_cache_block_offsets + # and is therefore max_blocks_per_seq long, so this + # means the sequence wants more blocks than + # max_seq_len allows -- i.e. it is consuming KV + # faster than it emits tokens. + raise ValueError( + f"User-provided base page indices is too short: " + f"buffer holds {len(indices)} blocks " + f"(max_blocks_per_seq) but the sequence now " + f"needs {new_num_blocks} (was {old_num_blocks}); " + f"num_blocks={self.num_blocks}, " + f"num_committed_tokens={self.num_committed_tokens}") stream_wait_events( self.cuda_stream, (s.ready_event for s in chain.from_iterable(slots)) diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 4f497787e143..159853c5f514 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1370,6 +1370,16 @@ class MemoryUpdateRequest(OpenAIBaseModel): tags: List[str] = Field(default=["model", "kv_cache"]) +class DSparkVerifyLenPinRequest(OpenAIBaseModel): + verify_len: Optional[int] = Field( + default=None, + description="DSpark verify length to pin every request to, or null to " + "clear the pin and hand scheduling back to the confidence planner. " + "Must be one of the captured verify-length tiers: an uncaptured " + "length drops every step out of CUDA-graph replay. Profiling knob -- " + "it bypasses the planner and is not a serving configuration.") + + class UpdateWeightsRequest(OpenAIBaseModel): weights: Optional[Dict[str, str]] = Field( default=None, diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 52dc5e488387..3c50959123e6 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -73,6 +73,7 @@ CompletionResponseChoice, EmbeddingRequest, EmbeddingResponse, EmbeddingResponseData, EmbeddingUsageInfo, ErrorResponse, ImageGenerationRequest, ImageGenerationResponse, ImageObject, + DSparkVerifyLenPinRequest, MemoryUpdateRequest, ModelCard, ModelList, PromptTokensDetails, ResponseFormat, ResponsesRequest, ResponsesResponse, TokenizeRequest, TokenizeResponse, UpdateWeightsRequest, UsageInfo, @@ -888,6 +889,9 @@ def register_routes(self): self.app.add_api_route("/server_info", self.get_server_info, methods=["GET"]) + self.app.add_api_route("/dspark/verify_len_pin", + self.set_dspark_verify_len_pin, + methods=["POST"]) if self.generator.args.return_perf_metrics: # register /prometheus/metrics self.mount_metrics() @@ -2381,6 +2385,44 @@ async def update_weights(self, args=(request.weights, )) return JSONResponse(content={"status": "success"}) + async def set_dspark_verify_len_pin( + self, request: DSparkVerifyLenPinRequest) -> JSONResponse: + """Pin (or clear) the DSpark verify length at runtime. + + Lets a cost-table sweep walk the verify-length ladder against a + RUNNING server -- the shape each cell needs in order to be labelled -- + instead of rebuilding the engine per length, which is what the + ``TLLM_DSPARK_FORCE_VERIFY_LEN`` environment variable requires and + which cannot reach ranks spawned before it was set. The pin takes + effect on the next decode step, simultaneously on every rank. + """ + # Dispatched by CAPABILITY, not by type. trtllm-serve builds a + # PyTorchLLM, which carries the synchronous `_collective_rpc`; only the + # AsyncLLM entry point has the awaitable `collective_rpc`. Asserting + # the type (as the memory endpoints do) rejects the very deployment + # this knob exists to profile, and does it with a 500. + rpc = getattr(self.generator, "collective_rpc", None) + sync_rpc = getattr(self.generator, "_collective_rpc", None) + if rpc is None and sync_rpc is None: + return JSONResponse( + content={"error": "this server exposes no worker RPC channel, " + "so the verify-length pin cannot be set"}, + status_code=400) + try: + if rpc is not None: + applied = await rpc('set_dspark_verify_len_pin', + args=(request.verify_len, )) + else: + applied = await asyncio.to_thread( + sync_rpc, 'set_dspark_verify_len_pin', + args=(request.verify_len, )) + except (ValueError, RuntimeError, NotImplementedError) as exc: + # A bad tier or a non-DSpark engine is a caller error, not a server + # fault: answer it rather than returning a 500 with a stack trace. + return JSONResponse(content={"error": str(exc)}, status_code=400) + value = applied[0] if isinstance(applied, list) and applied else applied + return JSONResponse(content={"verify_len": value}) + async def get_server_info(self) -> JSONResponse: content = {"disaggregated_params": self.generator.disaggregated_params} args = getattr(self.generator, "args", None) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f20e02169d62..73ad68b249ea 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1657,6 +1657,13 @@ "kind": "value", "path": "speculative_config.block_size" }, + { + "allowed_values": [], + "annotation": "Optional[List[Annotated[int, Gt(gt=0)]]]", + "converter": "", + "kind": "value", + "path": "speculative_config.confidence_verify_len_tiers" + }, { "allowed_values": [ "AUTO", @@ -1717,6 +1724,13 @@ "kind": "value", "path": "speculative_config.eagle_choices" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "speculative_config.enable_confidence_scheduling" + }, { "allowed_values": [], "annotation": "", @@ -1724,6 +1738,13 @@ "kind": "value", "path": "speculative_config.enable_global_pool" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "speculative_config.enable_ragged_verify" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index ad192d34d9a5..966c6f3e1082 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4001,6 +4001,187 @@ def test_gsm8k_dep8_megamoe_deepgemm(self): f"GSM8K accuracy {score:.3f} is below recorded reference " f"{acc_params.ref_accuracy:.3f}") + @pytest.mark.skip_less_mpi_world_size(8) + def test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling(self): + """Same accuracy bar with confidence-scheduled verification enabled. + + The scheduler only decides how MANY drafted tokens are sent to the + target; acceptance still goes through the unchanged target verify, so + the output distribution must be untouched. Any drop here means the + dynamic draft length broke an invariant (KV rewind amount, the draft + hidden-state gather stride, or the per-token sampling-parameter stride) + rather than that the scheduler made a bad throughput call -- those are + silent-wrong-answer bugs, which is why this asserts the SAME reference + accuracy as the baseline above rather than a relaxed one. + + Runs under attention-DP + TP8, so it also covers the cross-rank + agreement on the draft length: if ranks disagreed they would select + different CUDA graphs and the collectives would diverge. + """ + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.5) + spec_config = DSparkDecodingConfig(max_draft_len=5, + speculative_model=self.MODEL_PATH, + enable_confidence_scheduling=True) + with LLM(self.MODEL_PATH, + attn_backend="TRTLLM", + tensor_parallel_size=8, + moe_expert_parallel_size=8, + enable_attention_dp=True, + moe_config=MoeConfig(backend="MEGAMOE_DEEPGEMM"), + max_batch_size=DEEPSEEKV4_TEST_MAX_BATCH_SIZE, + max_seq_len=4096, + max_num_tokens=4096, + kv_cache_config=kv_cache_config, + enable_chunked_prefill=False, + disable_overlap_scheduler=True, + custom_tokenizer="deepseek_v4", + speculative_config=spec_config) as llm: + task = GSM8K(self.MODEL_NAME) + acc_params = task.get_hypothesis_testing_params( + dtype=llm.args.dtype, + quant_algo=llm.args.quant_config.quant_algo, + kv_cache_quant_algo=llm.args.quant_config.kv_cache_quant_algo, + spec_dec_algo=llm.args.speculative_config.decoding_type) + assert acc_params.num_samples == GSM8K.NUM_SAMPLES + with mock.patch.dict(os.environ, {"INTEGRATION_TEST": "0"}): + score = task.evaluate( + llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + assert score >= acc_params.ref_accuracy, ( + f"GSM8K accuracy {score:.3f} with confidence scheduling is below " + f"the recorded reference {acc_params.ref_accuracy:.3f}; the " + f"scheduler must not change the output distribution") + + def _run_gsm8k(self, llm): + task = GSM8K(self.MODEL_NAME) + acc_params = task.get_hypothesis_testing_params( + dtype=llm.args.dtype, + quant_algo=llm.args.quant_config.quant_algo, + kv_cache_quant_algo=llm.args.quant_config.kv_cache_quant_algo, + spec_dec_algo=llm.args.speculative_config.decoding_type) + assert acc_params.num_samples == GSM8K.NUM_SAMPLES + with mock.patch.dict(os.environ, {"INTEGRATION_TEST": "0"}): + score = task.evaluate( + llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + return score, acc_params + + @pytest.mark.skip_less_mpi_world_size(8) + @parametrize_with_ids("moe_backend", ["MEGAMOE_DEEPGEMM", "TRTLLM"]) + def test_gsm8k_dep8_static_verify_overlap(self, moe_backend): + """DSpark verifying the full block, with the overlap scheduler ON. + + Kept, rather than deleted, because it is the STATIC reference the other + two modes are differenced against: same prompts, same seed, verify + everything. cap-accept must match it token for token (the schedule only + changes how many tokens are committed per step, never which), and + compact must then match cap-accept. + + The sibling test above pins ``disable_overlap_scheduler=True``, which + makes it a weaker test than it looks: with overlap off, + ``_prepare_tp_inputs`` takes the branch that has no previous-batch + tensors, and the whole previous-step relay -- the KV-length correction, + the draft-token gather, the rewind -- is simply not exercised. Overlap + on is also the configuration anyone would actually serve. + + Cross-rank agreement is the other reason this matters. Under + attention-DP the draft length is part of the CUDA-graph key but is not + covered by the graph-eligibility allgather, so ranks that disagree pick + different graphs and their collectives diverge into a hang rather than a + wrong answer. + + Parameterized over both MoE backends. The scheduler changes how many + tokens reach the target, and the MoE chunks from that same count, so a + backend is a plausible place for a token-count assumption to hide even + though neither is supposed to know about speculation. + """ + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.5) + # No scheduling at all: every request verifies its whole drafted block. + # There is nothing for a cost table to inform here, which is the point + # -- this run has to be reproducible without one so it can serve as the + # reference for the two scheduled modes. + spec_config = DSparkDecodingConfig( + max_draft_len=5, + speculative_model=self.MODEL_PATH, + enable_confidence_scheduling=False) + with LLM(self.MODEL_PATH, + attn_backend="TRTLLM", + tensor_parallel_size=8, + moe_expert_parallel_size=8, + enable_attention_dp=True, + moe_config=MoeConfig(backend=moe_backend), + max_batch_size=DEEPSEEKV4_TEST_MAX_BATCH_SIZE, + max_seq_len=4096, + max_num_tokens=4096, + kv_cache_config=kv_cache_config, + cuda_graph_config=CudaGraphConfig(enable_padding=True), + enable_chunked_prefill=False, + disable_overlap_scheduler=False, + custom_tokenizer="deepseek_v4", + speculative_config=spec_config) as llm: + score, acc_params = self._run_gsm8k(llm) + assert score >= acc_params.ref_accuracy, ( + f"GSM8K accuracy {score:.3f} with confidence scheduling and the " + f"overlap scheduler is below the recorded reference " + f"{acc_params.ref_accuracy:.3f}") + + @pytest.mark.skip_less_mpi_world_size(8) + @parametrize_with_ids("moe_backend", ["MEGAMOE_DEEPGEMM", "TRTLLM"]) + def test_gsm8k_dep8_ragged_verify(self, moe_backend): + """Per-request verify windows, overlap scheduler on. + + ``enable_padding=True`` is required rather than incidental: the ragged + token budget is planned against the padded batch and reserves tokens for + the padding rows, so without them the step's total misses its captured + bucket. The config validator refuses the combination outright, which is + what this passes. + + Accuracy alone would not establish much here -- the planner declines to + trim without a profiled SPS cost table, and a run where it declined on + every step is a uniform baseline wearing a different config, scoring + exactly the same. A crafted non-flat cost table exists for that + reason: it hands out windows from the captured ladder deterministically, + which separates "is the ragged packing correct" from "does the planner + choose to trim". Set it in the environment when running this to test the + former. + """ + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.5) + # Without a cost table the planner is provably unable to trim: a flat + # table short-circuits `decide_verify_lens` before confidence is even + # read, so the ragged path stays dark and this test degenerates into the + # uniform baseline while still scoring the same. Point this at a table to + # exercise the real decision path. A table need not be *accurate* to be + # useful here -- scheduling only chooses how many drafted positions are + # submitted, and acceptance is unchanged, so an approximate table cannot + # change the output distribution, only the throughput it was picked for. + sps_table_path = os.environ.get("TLLM_DSPARK_SPS_TABLE") or None + spec_config = DSparkDecodingConfig( + max_draft_len=5, + speculative_model=self.MODEL_PATH, + enable_confidence_scheduling=True, + enable_ragged_verify=True, + confidence_sps_table_path=sps_table_path) + with LLM(self.MODEL_PATH, + attn_backend="TRTLLM", + tensor_parallel_size=8, + moe_expert_parallel_size=8, + enable_attention_dp=True, + moe_config=MoeConfig(backend=moe_backend), + max_batch_size=DEEPSEEKV4_TEST_MAX_BATCH_SIZE, + max_seq_len=4096, + max_num_tokens=4096, + kv_cache_config=kv_cache_config, + cuda_graph_config=CudaGraphConfig(enable_padding=True), + enable_chunked_prefill=False, + disable_overlap_scheduler=False, + custom_tokenizer="deepseek_v4", + speculative_config=spec_config) as llm: + score, acc_params = self._run_gsm8k(llm) + assert score >= acc_params.ref_accuracy, ( + f"GSM8K accuracy {score:.3f} with ragged verification is below " + f"the recorded reference {acc_params.ref_accuracy:.3f}") + @pytest.mark.timeout(14400) @pytest.mark.skip_less_device_memory(140000) diff --git a/tests/microbenchmarks/dspark_fit_sts.py b/tests/microbenchmarks/dspark_fit_sts.py new file mode 100644 index 000000000000..856b84326a9c --- /dev/null +++ b/tests/microbenchmarks/dspark_fit_sts.py @@ -0,0 +1,281 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fit the DSpark confidence head's per-position STS temperatures. + +WHAT IT FITS, AND WHY THAT OBJECTIVE +------------------------------------ +The planner never reads a per-position probability. It reads the *cumulative +product*:: + + survival[r][j] = prod_{i <= j} sigmoid(logit[r][i] / T[i]) + +and sums the largest of those to form tau, the numerator of the budget argmax. +So the thing that has to be calibrated is the survival, not the per-position +confidence -- a temperature vector that makes each position individually +well-calibrated can still leave the depth-5 product badly biased, because a +per-position error of x compounds to x^5. + +Hence the objective here is the expected calibration error **of the running +product** at each position, against the label "positions 0..j were ALL +accepted". Positions are fitted left to right and frozen, so every position is +optimised in the cumulative context it will actually be used in. + +This mirrors SGLang's ``benchmark/dspark_sts_fit.py`` (same 41-point log grid, +same left-to-right freeze, same cumprod-ECE objective), so a table fitted by +either tool is meaningful to the other. Deliberately so: the vectors are the one +artifact the two stacks can share without a converter. + +INPUT +----- +Shards written by ``DSparkStsRecorder`` (set ``TLLM_DSPARK_STS_COLLECT_PATH`` +on a normal DSpark run), each a torch save of:: + + {"logits": [N, K] float32, "prefix_mask": [N, K] float32} + +USAGE +----- + python3 tests/microbenchmarks/dspark_fit_sts.py \\ + --data "/path/to/sts_collect.*.pt" --out sts.json + +The emitted JSON carries both ``sts_temperatures`` and ``temperatures`` so it +loads in this repo and in SGLang unchanged. +""" + +import argparse +import glob +import json +import math +import sys +from typing import List, Tuple + +import torch + +_EPS_PROB = 1e-8 + + +def default_temperature_grid() -> torch.Tensor: + """41 points, log-spaced over [0.1, 10]. Same grid SGLang searches.""" + return torch.logspace(math.log10(0.1), math.log10(10.0), steps=41) + + +def expected_calibration_error(*, probs: torch.Tensor, targets: torch.Tensor, + num_bins: int) -> float: + """Bin by predicted probability; average |mean predicted - mean actual|. + + Weighted by bin population, so a bin holding three samples cannot dominate + one holding thirty thousand. + """ + probs = probs.reshape(-1).to(torch.float64).clamp(_EPS_PROB, 1.0 - _EPS_PROB) + targets = targets.reshape(-1).to(torch.float64) + total = probs.numel() + if total == 0: + return float("nan") + bin_index = (probs * num_bins).long().clamp_(0, num_bins - 1) + count = torch.zeros(num_bins, dtype=torch.float64) + pred_sum = torch.zeros(num_bins, dtype=torch.float64) + target_sum = torch.zeros(num_bins, dtype=torch.float64) + count.scatter_add_(0, bin_index, torch.ones_like(probs)) + pred_sum.scatter_add_(0, bin_index, probs) + target_sum.scatter_add_(0, bin_index, targets) + denom = count.clamp_min(1.0) + bin_error = (pred_sum / denom - target_sum / denom).abs() + return float((bin_error * count).sum().item() / total) + + +def fit_sts_temperatures(*, logits: torch.Tensor, prefix_mask: torch.Tensor, + grid: torch.Tensor, num_bins: int = 15) -> dict: + """Greedy left-to-right search minimising the cumprod ECE per position. + + Left to right and frozen, not joint: position j's temperature is chosen + against the survival the already-fitted positions 0..j-1 produce. A joint + search over K temperatures would be exponential in the grid, and the greedy + order is the causal one -- position j's contribution is only ever multiplied + into a product that already contains its predecessors. + """ + logits = logits.to(torch.float64) + prefix_mask = prefix_mask.to(torch.float64) + num_samples, block_size = logits.shape + if num_samples == 0: + raise ValueError("need at least one sample") + grid_values = grid.to(torch.float64).tolist() + + temperatures: List[float] = [] + ece_before: List[float] = [] + ece_after: List[float] = [] + + # Two running products: what an uncalibrated head would have said (T == 1), + # and what the fit says so far. Reporting both is the only way to see + # whether the fit bought anything. + survival_at_one = torch.ones(num_samples, dtype=torch.float64) + survival_fitted = torch.ones(num_samples, dtype=torch.float64) + + for position in range(block_size): + position_logits = logits[:, position] + position_target = prefix_mask[:, position] + + survival_at_one = survival_at_one * torch.sigmoid(position_logits) + ece_before.append( + expected_calibration_error(probs=survival_at_one, + targets=position_target, + num_bins=num_bins)) + + best_t = grid_values[0] + best_survival = survival_fitted * torch.sigmoid(position_logits / best_t) + best_ece = expected_calibration_error(probs=best_survival, + targets=position_target, + num_bins=num_bins) + for temperature in grid_values[1:]: + candidate = survival_fitted * torch.sigmoid(position_logits / temperature) + ece = expected_calibration_error(probs=candidate, + targets=position_target, + num_bins=num_bins) + if ece < best_ece: + best_ece, best_t, best_survival = ece, temperature, candidate + + temperatures.append(float(best_t)) + ece_after.append(float(best_ece)) + survival_fitted = best_survival + + return { + "temperatures": temperatures, + "ece_before": ece_before, + "ece_after": ece_after, + } + + +def load_shards(pattern: str) -> Tuple[torch.Tensor, torch.Tensor]: + """Load shards, refusing any whose pairing provenance is wrong or absent. + + Shards written by the current recorder carry ``meta.pairing == + "draft_seq_ring"``: each row's logits were selected from a stamped host + snapshot of exactly the draft pass that produced the block the label + verifies, keyed by the worker's own row allocator. Earlier recorders + paired the label against a sampler-time read of the live buffer -- stale + by one draft pass for most rows under the overlap scheduler (on job + 2562577, 26383 of 32083 rows differed, correlation 0.14) -- and keyed + rows by ``py_seq_slot``, a different allocator that drifts after the + first request completes. Fitting such a shard does not fail loudly; it + converges to a temperature vector that calibrates noise, which is what + the 0.23 residual ECE of the first deployed table was. There is + deliberately no flag to accept them: re-collect. + """ + paths = sorted(glob.glob(pattern)) + if not paths: + raise SystemExit(f"no shards matched {pattern!r}") + logits, masks, block_size = [], [], None + for path in paths: + blob = torch.load(path, map_location="cpu") + meta = blob.get("meta") or {} + pairing = meta.get("pairing") + if pairing != "draft_seq_ring": + raise SystemExit( + f"{path}: pairing provenance is {pairing!r}, expected " + f"'draft_seq_ring'. This shard predates the stamped snapshot " + f"ring, so its labels are joined to the wrong draft pass " + f"and/or the wrong buffer row, and a fit on it calibrates " + f"noise. Re-collect with the current recorder.") + lg, pm = blob["logits"], blob["prefix_mask"] + if lg.shape != pm.shape: + raise SystemExit(f"{path}: logits {tuple(lg.shape)} != mask {tuple(pm.shape)}") + if block_size is None: + block_size = lg.shape[1] + elif lg.shape[1] != block_size: + # Silently concatenating these would fit a vector for a block size + # that never existed. + raise SystemExit( + f"{path}: block size {lg.shape[1]} != {block_size} from earlier shards") + logits.append(lg) + masks.append(pm) + return torch.cat(logits, 0), torch.cat(masks, 0) + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--data", required=True, + help="glob for shards from DSparkStsRecorder") + ap.add_argument("--out", required=True, help="where to write the JSON") + ap.add_argument("--num-bins", type=int, default=15) + ap.add_argument("--min-samples", type=int, default=1000, + help="refuse to fit on fewer; a temperature fitted on a " + "handful of blocks is noise with a filename") + ap.add_argument("--eval-temps", default=None, + help="Evaluate an EXISTING temperature table on these " + "shards instead of fitting: reports the per-position " + "cumprod ECE the loaded vector achieves on fresh " + "data. This is the producer+consumer joint check -- " + "collect on a serving run with the table loaded and " + "the window pinned to the block, then compare this " + "number against the table's stored ece_after. A " + "large gap means the table does not describe the " + "live head (drift, wrong file, wrong column), and " + "nothing else in the pipeline can see that.") + args = ap.parse_args(argv) + + logits, prefix_mask = load_shards(args.data) + + if args.eval_temps: + with open(args.eval_temps, encoding="utf-8") as fh: + table = json.load(fh) + temps = table.get("sts_temperatures") or table["temperatures"] + lg = logits.to(torch.float64) + pm = prefix_mask.to(torch.float64) + survival = torch.ones(lg.shape[0], dtype=torch.float64) + print(f"evaluating {args.eval_temps} on {lg.shape[0]} fresh rows:") + stored = table.get("ece_after") or [] + for j in range(min(lg.shape[1], len(temps))): + survival = survival * torch.sigmoid(lg[:, j] / float(temps[j])) + ece = expected_calibration_error(probs=survival, + targets=pm[:, j], + num_bins=args.num_bins) + ref = f" (fit-time {stored[j]:.4f})" if j < len(stored) else "" + print(f" pos {j}: live ECE {ece:.4f}{ref}") + return 0 + n, block_size = logits.shape + print(f"loaded {n} samples, block_size={block_size}") + if n < args.min_samples: + raise SystemExit( + f"only {n} samples (< --min-samples {args.min_samples}). Collect " + f"more before fitting: the deepest positions are the sparsest and " + f"the ones the planner is most sensitive to.") + + # An all-ones column means every block was fully accepted at that depth, so + # the ECE is defined but carries no information to fit against -- say so + # rather than emitting a confident-looking temperature. + per_position_rate = prefix_mask.mean(dim=0) + print("prefix-accept rate by position:", + [round(float(v), 4) for v in per_position_rate]) + degenerate = [j for j, v in enumerate(per_position_rate) if v in (0.0, 1.0)] + if degenerate: + print(f"WARNING: positions {degenerate} are degenerate (rate 0 or 1); " + f"their temperature is unconstrained by this data") + + result = fit_sts_temperatures(logits=logits, prefix_mask=prefix_mask, + grid=default_temperature_grid(), + num_bins=args.num_bins) + + payload = { + "sts_temperatures": result["temperatures"], + "temperatures": result["temperatures"], + "ece_before": result["ece_before"], + "ece_after": result["ece_after"], + "dataset": args.data, + "num_samples": int(n), + } + with open(args.out, "w", encoding="utf-8") as fh: + json.dump(payload, fh, indent=2) + + print(f"\nwrote {args.out}") + print(f" temperatures : {[round(t, 4) for t in result['temperatures']]}") + print(f" ECE before : {[round(e, 5) for e in result['ece_before']]}") + print(f" ECE after : {[round(e, 5) for e in result['ece_after']]}") + improved = sum(1 for a, b in zip(result["ece_after"], result["ece_before"]) if a < b) + print(f" improved at {improved}/{block_size} positions") + if improved == 0: + print(" NOTE: no position improved -- the head was already calibrated " + "on this data, or the data does not cover the relevant range.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/microbenchmarks/dspark_make_steep_sps_table.py b/tests/microbenchmarks/dspark_make_steep_sps_table.py new file mode 100644 index 000000000000..077a78390306 --- /dev/null +++ b/tests/microbenchmarks/dspark_make_steep_sps_table.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Emit an SPS cost table steep enough that the planner chooses to trim. + +WHAT THIS IS FOR, AND WHAT IT IS NOT FOR +---------------------------------------- +It exercises the ragged path end to end *through the planner*: with a steep +``theta(M)`` the confidence scheduler decides to trim, and its budget then +flows through the same confidence-ordered top-k that ships, which hands +different windows to requests with different survival. That is the only way to +see per-request verify lengths on a checkpoint whose acceptance is too high for +trimming to pay. + +**No throughput or latency number taken with this table means anything.** The +planner is optimising against a fictional cost curve, so it will trim where a +real deployment should not. Use it to answer "does the ragged path correctly +support different verify lengths per request", never "is ragged faster". The +output carries ``"SYNTHETIC": true`` so a run cannot quietly present it as a +measurement -- ``run_dspark_throughput.sh`` refuses to treat such a table as +evidence. + +WHY A TABLE RATHER THAN AN OVERRIDE +----------------------------------- +A table drives the real confidence-scheduled path; it is the only sanctioned +mechanism for forcing window assignments. + +HOW STEEPNESS IS CHOSEN +----------------------- +The scheduler maximises expected accepted tokens per second, roughly + + (bs + sum of the top-K survival probabilities) / T(bs, K) + +With per-token acceptance around 0.9 each extra verify token adds about 0.9 to +the numerator, so trimming only wins once ``theta`` grows faster than that. The +default ``--exponent 2.0`` makes theta quadratic in the token count, which is +comfortably steeper than the measured curve (near-linear) and forces trimming +across the whole batch-size ladder. + + # derive from a real table, keeping its grid + python3 tests/microbenchmarks/dspark_make_steep_sps_table.py \ + --from sps_real_final.json --out steep.json + + # or synthesise a grid outright + python3 tests/microbenchmarks/dspark_make_steep_sps_table.py --out steep.json +""" + +import argparse +import json +import sys + + +def _default_grid(max_batch_size: int, width: int): + """Token totals a ragged batch can land on: bs * (verify_len + 1). + + The ladder always includes ``max_batch_size`` itself. A grid that stops + short of the deployment's widest step is worse than useless: theta clamps + at the last bucket, so the full window looks no more expensive than a + trimmed one and the planner keeps it -- a table that silently cannot + induce trimming, which is the failure this tool exists to prevent. + """ + ladder = sorted({b for b in (1, 4, 8, 16, 32, 64, 128, 256, 512) + if b <= max_batch_size} | {max_batch_size}) + tokens = sorted({b * w for b in ladder for w in range(2, width + 1)}) + return ladder, tokens + + + +def _preview_with_real_planner(*, path, batch_sizes, accept, max_verify_len): + """What budget_argmax_over_uniform_lens actually picks for this table. + + Returns True if any previewed batch size trims, False if none do, or None + when the planner cannot be imported (outside the container), so the caller + can distinguish "checked and fine" from "not checked". + """ + try: + import numpy as np + from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, budget_argmax_over_uniform_lens) + except Exception: + return None + + table = SpsCostTable.from_json_file(path) if hasattr( + SpsCostTable, "from_json_file") else None + if table is None: + with open(path, encoding="utf-8") as handle: + table = SpsCostTable(**{ + k: v for k, v in json.load(handle).items() + if not k.startswith("_") and k != "SYNTHETIC" + }) + + tiers = sorted({1, max(1, (max_verify_len + 1) // 2), max_verify_len}) + print() + print(f"planner's actual choice at acceptance {accept} " + f"(L = verify length, {max_verify_len} = no trim), tiers {tiers}:") + any_trim = False + for bs in [int(b) for b in str(batch_sizes).split(",") if str(b).strip()]: + # survival[r, k] = P(first k+1 drafted tokens all accepted) + survival = np.array( + [[accept**(k + 1) for k in range(max_verify_len)]] * bs, + dtype=np.float64) + chosen = budget_argmax_over_uniform_lens( + survival=survival, num_gen_requests=bs, cost_table=table, + allowed_lens=tiers, min_verify_len=1) + trims = int(chosen) < max_verify_len + any_trim = any_trim or trims + print(f" bs={bs:<5} -> L={chosen}" + f"{' TRIM' if trims else ' (full window)'}") + return any_trim + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--out", required=True) + parser.add_argument("--from", dest="source", default=None, + help="real table to take the grid from; its timings " + "are replaced, only token_counts/batch_sizes are " + "reused so the shape matches the deployment") + parser.add_argument("--exponent", type=float, default=2.0, + help="theta(M) ~ M**exponent; >1 makes extra verify " + "tokens superlinearly expensive") + parser.add_argument("--peak-ms", type=float, default=400.0, + help="theta at the largest token count") + parser.add_argument("--fixed-overhead-ms", type=float, default=0.05, + help="the non-trimmable floor. MUST be small next to " + "theta in the operating range, or trimming cannot " + "pay however steep theta is -- see --preview-bs") + parser.add_argument("--preview-bs", default="8,32,128", + help="batch sizes to predict the planner's choice for, " + "so a table that cannot induce trimming is caught " + "here instead of 20 minutes into a run") + parser.add_argument("--preview-accept", type=float, default=0.9, + help="per-token acceptance to predict against; 0.9 is " + "what DeepSeek-V4-Pro-DSpark measures") + parser.add_argument("--max-batch-size", type=int, default=128) + parser.add_argument("--max-verify-len", type=int, default=5) + args = parser.parse_args() + + if args.exponent <= 1.0: + parser.error( + f"--exponent {args.exponent} is not steeper than linear, so the " + f"planner will keep verifying the full block and the ragged path " + f"will stay dark -- which is the exact failure this tool exists " + f"to avoid") + + needed = args.max_batch_size * (args.max_verify_len + 1) + if args.source: + with open(args.source, encoding="utf-8") as handle: + src = json.load(handle) + tokens = [int(t) for t in src["token_counts"]] + batch_sizes = [int(b) for b in src["batch_sizes"]] + if max(tokens) < needed: + parser.error( + f"--from {args.source} tops out at {max(tokens)} tokens but " + f"max_batch_size {args.max_batch_size} reaches {needed} at the " + f"full window. theta would clamp there, making the full window " + f"look as cheap as a trimmed one, so the planner would never " + f"trim. Drop --from to synthesise a grid that covers it.") + else: + batch_sizes, tokens = _default_grid(args.max_batch_size, + args.max_verify_len + 1) + + tokens = sorted({t for t in tokens if t > 0}) + peak = float(max(tokens)) + theta = [args.peak_ms * (t / peak)**args.exponent for t in tokens] + + table = { + "token_counts": tokens, + "step_time_ms": theta, + "fixed_overhead_ms": args.fixed_overhead_ms, + "batch_sizes": batch_sizes, + "batch_overhead_ms": [0.0] * len(batch_sizes), + "SYNTHETIC": True, + "_meta": { + "purpose": "force the planner to trim so the ragged path runs; " + "NOT a measurement of anything", + "exponent": args.exponent, + "derived_from": args.source, + }, + } + with open(args.out, "w", encoding="utf-8") as handle: + json.dump(table, handle, indent=2) + handle.write("\n") + + # Ask the REAL planner what it will do, rather than re-deriving its argmax. + # + # A hand-written approximation of the objective was wrong twice: it models + # the cost as fixed + theta(M), while budget_argmax_over_uniform_lens calls + # cost_table.step_time(total_verify_tokens(bs, length), num_gen_requests), + # which interpolates the table and adds the per-batch term. The + # approximation predicted "trim at every batch size" for a table under + # which the planner kept the full window through a 20-minute run. A preview + # that does not agree with the thing it previews is worse than none. + any_trim = _preview_with_real_planner( + path=args.out, batch_sizes=args.preview_bs, + accept=args.preview_accept, max_verify_len=args.max_verify_len) + if any_trim is None: # planner not importable here + print("preview skipped: run this inside the TensorRT-LLM container to " + "have the real planner check the table before you spend a run " + "on it", file=sys.stderr) + any_trim = True + if not any_trim: + print() + print("WARNING: this table induces no trimming at any previewed batch " + "size, so the ragged path will stay dark. Lower " + "--fixed-overhead-ms or raise --exponent.", file=sys.stderr) + + print() + print(f"wrote {args.out}: {len(tokens)} token buckets " + f"{tokens[0]}..{tokens[-1]}, theta {theta[0]:.3f}..{theta[-1]:.1f} ms " + f"(exponent {args.exponent}), SYNTHETIC=true") + print("Use only to check that per-request verify lengths work. Any " + "throughput taken with this table is meaningless.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/microbenchmarks/dspark_sps_profiler.py b/tests/microbenchmarks/dspark_sps_profiler.py new file mode 100644 index 000000000000..4500b666b4b3 --- /dev/null +++ b/tests/microbenchmarks/dspark_sps_profiler.py @@ -0,0 +1,2633 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Measure the DSpark SPS cost table this deployment actually has. + +Why this exists +--------------- +:class:`~.dspark_planner.SpsCostTable` is the only thing that tells the verify +planner a long verification costs more than a short one. Without a measured +table ``SpsCostTable.flat()`` is installed, ``is_flat`` is true, and +``DSparkVerifyPlanner.decide_verify_lens`` returns ``None`` on every step +(``stats["fallback_flat_cost"]``) -- ragged verification is then a no-op that +still produces correct output at baseline accuracy, i.e. it fails silently. +A crafted (non-flat) cost table unblocks *correctness* testing by making the +planner itself trim, rather than bypassing +that gate, but it rotates a fixed ladder and knows nothing about cost, so it +cannot answer "is trimming profitable here". Only a measured table can. + +Run it:: + + python tests/microbenchmarks/dspark_sps_profiler.py \\ + --model /models/DeepSeek-V4-Pro-DSpark --tp-size 8 --ep-size 8 \\ + --enable-attention-dp --batch-sizes 8,16,32,64,128 \\ + --max-draft-len 5 --out sps_table.json + +and then point the serving config at the result:: + + speculative_config: + decoding_type: DSpark + enable_confidence_scheduling: true + confidence_sps_table_path: sps_table.json + confidence_verify_len_tiers: [...] # from _meta, see "Tiers" below + +What the sweep measures, and how the three terms are separated +-------------------------------------------------------------- +The planner's cost model is additive:: + + T(bs, M) = bias + alpha(bs) + theta(M), M = bs * (verify_len + 1) + +``theta(M)`` is the *only* term trimming reduces; ``bias`` and ``alpha(bs)`` are +paid whatever the planner picks. Because the planner maximizes ``tau / T``, a +**ratio**, the non-trimmable terms do not cancel: drop them and every verified +token looks proportionally more expensive than it is, and the planner +over-trims. So all three have to be measured, not guessed. + +The sweep is a cross product of batch size ``bs`` and *uniform* verify length +``L``. Each cell is one whole steady-state decode step time ``S(bs, L)``, and +sits at ``M = bs * (L + 1)`` (:func:`~.dspark_planner.total_verify_tokens` -- +the bonus token is included; getting that wrong is a whole ``bs`` of error, +easily a shelf's width). The cross product is what makes the terms separable: + +* **theta(M) is isolated along a row.** Fix ``bs`` and move ``L``: ``bias`` and + ``alpha(bs)`` are identical for every cell in the row, so the *differences* + along it are pure ``theta``. +* **alpha(bs) is isolated along a column.** The grid supplies same-``M``, + different-``bs`` collisions for free -- ``(bs, L)`` and ``(2*bs, L')`` land on + the same ``M`` whenever ``L + 1 == 2 * (L' + 1)``, e.g. ``(16, 3)`` and + ``(32, 1)`` both give ``M = 64``. Across such a collision ``theta`` cancels + and the difference is pure ``alpha``. + This is also why the batch-size ladder must be *connected* through those + collisions; :func:`fit_additive_cost_model` refuses a sweep that is not, + because each disconnected component would carry its own arbitrary offset and + the emitted ``theta`` staircase would splice two different scales together. +* **The remaining global constant is fixed by extrapolating theta to M -> 0.** + An additive fit only determines ``theta`` and ``alpha`` up to one shared + constant. We shift so that ``theta`` linearly extrapolates to zero at + ``M = 0``, which is exactly the definition that makes ``bias + alpha(bs)`` + the step's ``M``-independent floor -- the quantity the planner needs. +* **bias vs alpha(bs) is then a presentation split.** Only their sum enters + :meth:`SpsCostTable.step_time`, so we report ``bias = min_bs (bias + + alpha(bs))`` -- the floor present even at the smallest profiled batch -- and + give the rest to ``alpha``. Both are emitted so that a reader can see which + part is batch-driven. + +The emitted table is therefore in the **decomposed** encoding: ``step_time_ms`` +is ``theta(M)`` *only*. Do not also read it as whole-step time and re-add the +overheads -- that double-counts the floor, inflates ``T``, flattens the relative +differences between candidate lengths and drags the argmax back to verify-all. + +Holding verification uniform +---------------------------- +The run pins ``TLLM_DSPARK_RAGGED_VERIFY_MODE=static`` and clears +a crafted cost table. Under ``static`` the planner still chooses a +single batch-wide length, but the cost table it sees is the flat default, whose +``is_flat`` gate short-circuits to ``max_tier`` *before* any confidence is read +-- so the length is deterministically ``max_draft_len`` on every step and +``M = bs * (max_draft_len + 1)`` exactly. If the scheduler were allowed to trim +mid-sweep, each cell would be measuring a different ``M`` than the one it is +filed under, i.e. a moving target. + +That determinism is also why the ``L`` axis costs one engine build per value: +nothing in the runtime pins a *uniform* verify length below the full block +(a ragged split would break the M = bs*(L+1) filing, and there is no +analogue of SGLang's ``dspark_force_budget_frac``), so ``L`` is swept by +rebuilding with ``max_draft_len = block_size = L``. Use ``--samples-out`` / +``--fit-only`` to split a long sweep across jobs and refit without a GPU. + +**Known bias of that choice:** shrinking the block also shrinks the *draft* +pass, whose cost is roughly linear in the block length, while at deployment the +block is always drafted in full and only verification is trimmed. The draft +model is a few layers against a full-size target, so the leak is a small +fraction of the measured slope -- but it is a positive fraction, so the emitted +``theta`` slope is mildly overstated and the planner is biased towards +*over*-trimming. Prefer profiling only the lengths the ladder can actually +select, and treat a marginal trim decision as noise. + +Acceptance is pinned +-------------------- +``--pin-acceptance`` (on by default) sets +``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` to a tiny positive value, which +makes every request commit exactly its bonus token every step. ``0.0`` means +"override disabled", hence the epsilon rather than a clean zero. This matters +because acceptance drives how fast the KV cache grows, and KV length drives +attention cost: without pinning, a long-block cell would accumulate context +faster than a short-block one and the difference would be misattributed to +``theta(M)``. + +Statistics +---------- +Only *aligned* steps count: a step where every attention-DP rank reports the +same generation-request count and no context requests. The first +``--warmup-steps`` aligned steps of each cell are discarded (allocator growth, +graph warmup, autotuner) and the cell reports the **median** of the rest -- +never the mean. Decode step time has a long right tail from allocator and +scheduling noise, and a mean turns that tail into a fake riser, which is exactly +the failure a staircase cost model amplifies. + +Refusals +-------- +The tool would rather emit nothing than emit a table that looks profiled and +is not: + +* :class:`FlatCostTableError` -- the curve does not move, so ``is_flat`` is true + (or the total step time varies by less than the noise floor). The planner + would treat the file as unprofiled and keep verifying everything. +* :class:`InertCostTableError` -- the curve moves, but not between any two tiers + the planner may select at any profiled batch size, so the argmax can never + leave ``max_tier``. That is a non-flat table that trims nothing, and unlike a + flat one it does *not* trip ``fallback_flat_cost``, so it would look healthy. + +Tiers +----- +``_meta.recommended_confidence_verify_len_tiers`` is derived from the measured +shelves at the modal batch size, but **nothing reads it from this file**: +``DSparkDecodingConfig.verify_len_tiers`` always returns a non-empty list, so +``dspark.py``'s ``derive_verify_len_tiers`` fallback never runs. Copy the value +into ``confidence_verify_len_tiers`` by hand. The config is the right single +source anyway -- ``_get_graphs_to_capture`` and ``ragged_verify_token_buckets`` +read the same property, so a ladder that lives only here would let the planner +pick lengths that have no captured CUDA graph and silently drop the step into +eager execution. +""" + +from __future__ import annotations + +import argparse +import json +import os +import pathlib +import random +import re +import statistics +import sys +import threading +import time +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Sequence, Tuple + +import numpy as np + +from tensorrt_llm._torch.speculative.dspark_observability import RAGGED_VERIFY_MODE_ENV, RaggedVerifyMode +from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, + budget_argmax_over_uniform_lens, + derive_verify_len_tiers, + total_verify_tokens, +) + +__all__ = [ + "AdditiveFit", + "CellStat", + "FlatCostTableError", + "InertCostTableError", + "InsufficientSamplesError", + "ProfilerError", + "StepSample", + "SweepConfig", + "SweepGeometryError", + "aligned_steps_from_stats", + "build_cost_table_payload", + "check_table_is_informative", + "compress_to_risers", + "fit_additive_cost_model", + "load_cost_table", + "main", + "profitability_probe", + "run_sweep", + "running_max", + "summarize_cells", +] + +#: Per-iteration timing fields published by ``PyExecutor._profiler``. The host +#: field is the default because it is a clean single-loop CPU wall on every +#: scheduler; ``iterLatencyMS`` is deliberately not offered, since under the +#: overlap scheduler (the shipped default) it spans ~2 loops. +HOST_STEP_TIME_KEY = "hostStepTimeMS" +GPU_FORWARD_TIME_KEY = "gpuForwardTimeMS" +TIMING_FIELDS = {"host": HOST_STEP_TIME_KEY, "gpu_forward": GPU_FORWARD_TIME_KEY} + +#: ``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` treats 0.0 as "override off", +#: so pinning to bonus-token-only needs a value that floors to 0 without being 0. +#: The fractional part is the probability of one extra token, so 1e-6 fires +#: roughly never while still selecting the override path. +BONUS_TOKEN_ONLY_ACCEPTANCE = 1e-6 + +#: ``SpsCostTable.__post_init__`` rejects a non-positive ``step_time_ms``. After +#: normalizing theta to hit zero at M = 0 the smallest shelf can land at or below +#: zero, so it is clamped to this instead. +MIN_THETA_MS = 1e-4 + +#: How far the additive model may miss a cell, relative to a typical step, +#: before the fit says so. Not a repair threshold -- see +#: :func:`fit_additive_cost_model` for why the misfit is reported, not fixed. +DEFAULT_ADDITIVE_TOLERANCE = 0.02 + + +class ProfilerError(RuntimeError): + """Base class for every refusal this module raises.""" + + +class InsufficientSamplesError(ProfilerError): + """A cell did not produce enough aligned steady-state steps to trust.""" + + +class SweepGeometryError(ProfilerError): + """The measured cells cannot separate ``alpha(bs)`` from ``theta(M)``.""" + + +class FlatCostTableError(ProfilerError): + """The measured curve is flat, so the table would read as unprofiled.""" + + +class InertCostTableError(ProfilerError): + """The curve moves, but never between two selectable verify lengths.""" + + +# --------------------------------------------------------------------------- +# raw samples and per-cell statistics +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class StepSample: + """One aligned decode step, filed under the cell that produced it.""" + + batch_size: int + verify_len: int + step_time_ms: float + iteration: int = -1 + #: Ragged (frac-sweep) steps: the group's executed token total, read from + #: the stats rows. None on pinned sweeps, where M is derived instead. + measured_tokens: Optional[int] = None + + @property + def total_verify_tokens(self) -> int: + """``M`` for this step. Measured when the shape was ragged; derived + from the pin otherwise, so it cannot drift from the label.""" + if self.measured_tokens is not None: + return int(self.measured_tokens) + return total_verify_tokens(self.batch_size, self.verify_len) + + def to_json(self) -> dict: + payload = { + "batch_size": int(self.batch_size), + "verify_len": int(self.verify_len), + "step_time_ms": float(self.step_time_ms), + "iteration": int(self.iteration), + } + if self.measured_tokens is not None: + payload["measured_tokens"] = int(self.measured_tokens) + return payload + + @classmethod + def from_json(cls, raw: dict) -> "StepSample": + measured = raw.get("measured_tokens") + return cls( + batch_size=int(raw["batch_size"]), + verify_len=int(raw["verify_len"]), + step_time_ms=float(raw["step_time_ms"]), + iteration=int(raw.get("iteration", -1)), + measured_tokens=None if measured is None else int(measured), + ) + + +@dataclass(frozen=True) +class CellStat: + """The summarized cost of one ``(batch_size, verify_len)`` cell.""" + + batch_size: int + verify_len: int + step_time_ms: float + num_samples: int + p10_ms: float + p90_ms: float + #: Ragged (frac-sweep) cells: the executed token total the samples agreed + #: on. None on pinned sweeps. + measured_tokens: Optional[int] = None + + @property + def total_verify_tokens(self) -> int: + if self.measured_tokens is not None: + return int(self.measured_tokens) + return total_verify_tokens(self.batch_size, self.verify_len) + + def to_json(self) -> dict: + return { + "batch_size": int(self.batch_size), + "verify_len": int(self.verify_len), + "total_verify_tokens": int(self.total_verify_tokens), + "step_time_ms": round(float(self.step_time_ms), 6), + "num_samples": int(self.num_samples), + "p10_ms": round(float(self.p10_ms), 6), + "p90_ms": round(float(self.p90_ms), 6), + } + + +def aligned_steps_from_stats( + rows: Sequence[dict], + *, + expected_ranks: int = 1, + timing_key: str = HOST_STEP_TIME_KEY, +) -> List[Tuple[int, int, float]]: + """Pull ``(iteration, num_gen_requests, step_time_ms)`` out of ``get_stats``. + + A step is kept only when it is *aligned*: every expected attention-DP rank + published a row for that iteration, no rank had context requests in flight, + and all ranks scheduled the same number of generation requests. Anything + else is a step whose shape is not the cell's shape -- a prefill iteration, a + ramp step, a rank that fell behind -- and averaging it in is how a cost + table acquires risers that are not there. + + The timing is taken from the rank-0 row rather than averaged across ranks: + under attention-DP :mod:`~..pyexecutor.adp_iter_stats` fans out one record + per rank with *rank-local counters* but *rank-0's* step time, so the + per-rank timings are copies, not independent measurements. + """ + by_iteration: Dict[int, List[dict]] = {} + for row in rows: + iteration = row.get("iter") + if iteration is None: + continue + by_iteration.setdefault(int(iteration), []).append(row) + + aligned: List[Tuple[int, int, float]] = [] + for iteration in sorted(by_iteration): + group = sorted(by_iteration[iteration], key=lambda r: int(r.get("attentionDpRank", 0))) + if len({int(r.get("attentionDpRank", 0)) for r in group}) != int(expected_ranks): + continue + batching = [row.get("inflightBatchingStats") or {} for row in group] + if any(int(stats.get("numContextRequests", 0) or 0) for stats in batching): + continue + gen_counts = {int(stats.get("numGenRequests", 0) or 0) for stats in batching} + if len(gen_counts) != 1: + continue + num_gen_requests = gen_counts.pop() + if num_gen_requests <= 0: + continue + # gpuForwardTimeMS is read without a device sync, so it is absent on the + # steps whose event pair had not landed yet; those steps are simply not + # usable for that field. + timing = group[0].get(timing_key) + if timing is None or float(timing) <= 0.0: + continue + aligned.append((iteration, num_gen_requests, float(timing))) + return aligned + + +def ragged_steps_from_stats( + rows: Sequence[dict], + *, + expected_ranks: int = 1, + timing_key: str = HOST_STEP_TIME_KEY, +) -> List[Tuple[int, int, int, float]]: + """``(iteration, num_gen_requests, executed_tokens, step_time_ms)`` steps. + + The ragged (frac-sweep) sibling of :func:`aligned_steps_from_stats`: on top + of its alignment rules, every rank must also report the same + ``numGenerationTokens``. Under attention-DP the executed token total is the + group's agreed bucket -- every rank pads and fills to the same shape -- so + a step whose ranks disagree is a ramp/fallback step whose shape is not one + shape at all, and averaging it in would file its cost under a bucket it did + not run. + """ + by_iteration: Dict[int, List[dict]] = {} + for row in rows: + iteration = row.get("iter") + if iteration is None: + continue + by_iteration.setdefault(int(iteration), []).append(row) + + kept: List[Tuple[int, int, int, float]] = [] + for iteration in sorted(by_iteration): + group = by_iteration[iteration] + if len({int(r.get("attentionDpRank", 0)) for r in group}) != int(expected_ranks): + continue + batching = [row.get("inflightBatchingStats") or {} for row in group] + if any(int(stats.get("numContextRequests", 0) or 0) for stats in batching): + continue + gen_counts = {int(stats.get("numGenRequests", 0) or 0) for stats in batching} + if len(gen_counts) != 1 or (num_gen := gen_counts.pop()) <= 0: + continue + token_counts = {row.get("numGenerationTokens") for row in group} + if len(token_counts) != 1: + continue + tokens = token_counts.pop() + if tokens is None or int(tokens) <= 0: + continue + timing = group[0].get(timing_key) + if timing is None or float(timing) <= 0.0: + continue + kept.append((iteration, num_gen, int(tokens), float(timing))) + return kept + + +def explain_alignment( + rows: Sequence[dict], + *, + expected_ranks: int = 1, + timing_key: str = HOST_STEP_TIME_KEY, +) -> str: + """Say why :func:`aligned_steps_from_stats` kept nothing. + + A cell that reports ``aligned_steps=0`` is otherwise indistinguishable from + a cell that never ran, and the four rejection reasons want opposite fixes: + a rank count that is never right means ``--enable-attention-dp`` disagrees + with the deployment, context requests on every step mean the prompts are + still being ingested, disagreeing generation counts mean the batch is + draining, and a missing timing field means the requested ``--timing-field`` + is not being published. Guessing between those costs a full sweep each + time. + """ + by_iteration: Dict[int, List[dict]] = {} + for row in rows: + iteration = row.get("iter") + if iteration is None: + continue + by_iteration.setdefault(int(iteration), []).append(row) + if not by_iteration: + return (f"{len(rows)} stats rows carried no 'iter' field at all" + if rows else "no stats rows were received at all") + + reasons: Dict[str, int] = {} + rank_counts: Dict[int, int] = {} + gen_seen: set = set() + for iteration, group in by_iteration.items(): + ranks = {int(r.get("attentionDpRank", 0)) for r in group} + rank_counts[len(ranks)] = rank_counts.get(len(ranks), 0) + 1 + if len(ranks) != int(expected_ranks): + reasons["rank count != expected"] = reasons.get("rank count != expected", 0) + 1 + continue + batching = [row.get("inflightBatchingStats") or {} for row in group] + if any(int(stats.get("numContextRequests", 0) or 0) for stats in batching): + reasons["context requests in flight"] = reasons.get("context requests in flight", 0) + 1 + continue + gen_counts = {int(stats.get("numGenRequests", 0) or 0) for stats in batching} + gen_seen.update(gen_counts) + if len(gen_counts) != 1: + reasons["ranks disagree on numGenRequests"] = reasons.get("ranks disagree on numGenRequests", 0) + 1 + continue + if gen_counts and max(gen_counts) <= 0: + reasons["no generation requests"] = reasons.get("no generation requests", 0) + 1 + continue + timing = group[0].get(timing_key) + if timing is None or float(timing) <= 0.0: + reasons[f"missing/zero {timing_key}"] = reasons.get(f"missing/zero {timing_key}", 0) + 1 + continue + reasons["kept"] = reasons.get("kept", 0) + 1 + + parts = [f"{len(rows)} rows over {len(by_iteration)} iterations " + f"(iter {min(by_iteration)}..{max(by_iteration)})"] + parts.append("rejections: " + ", ".join( + f"{reason} x{count}" for reason, count in + sorted(reasons.items(), key=lambda kv: -kv[1])) or "none") + parts.append("ranks-per-iteration histogram: " + ", ".join( + f"{n}->{c}" for n, c in sorted(rank_counts.items())) + + f" (expected {expected_ranks})") + if gen_seen: + parts.append(f"numGenRequests observed: {sorted(gen_seen)[:12]}") + return "; ".join(parts) + + +def summarize_cells( + samples: Sequence[StepSample], + *, + warmup_steps: int, + min_samples: int, + expected_cells: Optional[Sequence[Tuple[int, int]]] = None, +) -> List[CellStat]: + """Median per ``(batch_size, verify_len)`` cell, after dropping warmup. + + Median, not mean: a decode step's distribution is a tight body with a long + right tail (allocator growth, host scheduling hiccups, the occasional + stats/log flush). The mean tracks the tail, and a staircase cost model reads + a tail-inflated point as a riser -- i.e. as a place worth trimming before. + + Args: + warmup_steps: leading aligned steps to discard *per cell*. They are kept + in the raw sample file so a refit can choose a different cut. + min_samples: post-warmup samples a cell must have to be admitted. + expected_cells: the requested grid, when there is one. Sources that + observe a deployment (a served sweep's iteration log, a production + log) inevitably capture ramp and drain steps whose occupancy was + never requested -- e.g. a batch of 2 while the client's 256 + requests were still being admitted. Those incidental cells are + dropped (and reported) when thin; only a thin cell that was + actually REQUESTED is an error. None means every observed cell was + asked for (the engine-building sweep, existing callers). + + Raises: + InsufficientSamplesError: if a requested cell falls short. Silently + dropping one of those would punch a hole in the grid and can + disconnect the batch-size ladder, which changes the fit rather + than just widening its error bars. + """ + grouped: Dict[Tuple[int, int, Optional[int]], List[StepSample]] = {} + for sample in samples: + key = (int(sample.batch_size), int(sample.verify_len), + None if sample.measured_tokens is None else int(sample.measured_tokens)) + grouped.setdefault(key, []).append(sample) + + requested = None if expected_cells is None else { + (int(b), int(l)) for b, l in expected_cells} + dropped: List[str] = [] + cells: List[CellStat] = [] + for (batch_size, verify_len, measured_tokens), cell_samples in sorted( + grouped.items(), key=lambda kv: (kv[0][0], kv[0][1], kv[0][2] or -1)): + ordered = sorted(cell_samples, key=lambda s: s.iteration) + steady = [s.step_time_ms for s in ordered[int(warmup_steps) :]] + if len(steady) < int(min_samples): + # Ragged sweeps visit whichever buckets the planner picked, so a + # thin incidental bucket is normal coverage noise, not a failed + # cell: drop it loudly rather than aborting the fit. + if measured_tokens is not None: + dropped.append(f"(bs={batch_size}, M={measured_tokens}, n={len(steady)})") + continue + if requested is not None and (batch_size, verify_len) not in requested: + dropped.append(f"(bs={batch_size}, L={verify_len}, n={len(steady)})") + continue + raise InsufficientSamplesError( + f"cell (batch_size={batch_size}, verify_len={verify_len}) has " + f"{len(steady)} steady samples after discarding " + f"{warmup_steps} warmup steps, need {min_samples}. Either the " + f"cell never reached a steady batch of {batch_size} generating " + f"requests (check that the batch size is captured by the CUDA " + f"graph config and that enough requests were in flight), or the " + f"measurement window was too short -- raise --measure-steps." + ) + values = np.asarray(steady, dtype=np.float64) + cells.append( + CellStat( + batch_size=batch_size, + verify_len=verify_len, + step_time_ms=float(np.median(values)), + num_samples=int(values.size), + p10_ms=float(np.percentile(values, 10)), + p90_ms=float(np.percentile(values, 90)), + measured_tokens=measured_tokens, + ) + ) + if dropped: + # Reported, not silent: a fit that quietly ignored data reads as + # "covered everything" when it did not. + print(f"[dspark-sps] dropped {len(dropped)} thin incidental cell(s) " + f"outside the requested grid: {', '.join(dropped[:12])}" + + (" ..." if len(dropped) > 12 else ""), file=sys.stderr) + return cells + + +# --------------------------------------------------------------------------- +# additive fit: T(bs, M) = bias + alpha(bs) + theta(M) +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class AdditiveFit: + """The separated cost terms, before they are shaped into a staircase.""" + + #: ``bias + alpha(bs)`` per profiled batch size, normalized so that theta + #: extrapolates to zero at ``M = 0``. + intercept_ms: Dict[int, float] + #: ``theta(M)`` on the measured token counts, non-decreasing. + theta_ms: Dict[int, float] + #: Largest ``|S - intercept - theta|`` over the cells, absolute and relative + #: to the median cell time. This is the additive model's own error bar: if + #: it is large, ``theta`` is not a function of ``M`` alone and the + #: decomposed encoding is not safe at batch sizes far from those profiled. + max_abs_residual_ms: float + max_rel_residual: float + warnings: Tuple[str, ...] = () + + +def _assert_grid_is_connected(cells: Sequence[CellStat]) -> None: + """Refuse a sweep whose ``(bs, M)`` bipartite graph has several components. + + Within one component the additive split is determined up to a single shared + constant, which the caller fixes. Across components each has its *own* + unknown constant, so concatenating their ``theta`` values into one staircase + silently splices two different scales -- the resulting risers are artifacts + of the offset, not of the hardware. + """ + parent: Dict[Tuple[str, int], Tuple[str, int]] = {} + + def find(node): + parent.setdefault(node, node) + while parent[node] != node: + parent[node] = parent[parent[node]] + node = parent[node] + return node + + for cell in cells: + a, b = find(("bs", cell.batch_size)), find(("m", cell.total_verify_tokens)) + if a != b: + parent[a] = b + + components = {find(("bs", cell.batch_size)) for cell in cells} + if len(components) > 1: + groups: Dict[Tuple[str, int], List[int]] = {} + for cell in cells: + groups.setdefault(find(("bs", cell.batch_size)), []).append(cell.batch_size) + islands = sorted(sorted(set(v)) for v in groups.values()) + raise SweepGeometryError( + f"the swept cells split into {len(components)} disconnected groups of " + f"batch sizes {islands}: no two groups share a total-token count, so " + f"alpha(bs) and theta(M) cannot be put on a common scale. Two batch " + f"sizes meet when bs1*(L1+1) == bs2*(L2+1) for swept lengths, e.g. a " + f"doubling ladder (8,16,32,...) with verify lengths spanning a factor " + f"of two (1..5 gives token/request ratios 2..6). Add the missing " + f"batch sizes or widen --verify-lens." + ) + + +def running_max(values: Sequence[float]) -> List[float]: + """Cheapest possible isotonic fit: enforce non-decreasing in place. + + ``SpsCostTable`` does not check monotonicity, and a noise-induced dip makes + a *longer* verification look cheaper than a shorter one -- the planner will + take that offer. Clamping up rather than down is the conservative direction: + it can only make the planner trim less. + """ + out: List[float] = [] + highest = -float("inf") + for value in values: + highest = max(highest, float(value)) + out.append(highest) + return out + + +def _least_squares_intercept(x: Sequence[float], y: Sequence[float]) -> float: + """Value of the best-fit line ``y = a + b*x`` at ``x = 0``.""" + xs = np.asarray(x, dtype=np.float64) + ys = np.asarray(y, dtype=np.float64) + if xs.size < 2 or np.ptp(xs) == 0.0: + return float(ys.min()) + slope = float(np.cov(xs, ys, bias=True)[0, 1] / np.var(xs)) + return float(ys.mean() - slope * xs.mean()) + + +def _solve_additive( + cells: Sequence[CellStat], +) -> Tuple[Dict[int, float], Dict[int, float], np.ndarray]: + """Least-squares solve of ``S = intercept[bs] + theta[M]`` over the cells. + + The design is rank-deficient by exactly one -- adding a constant to every + ``theta`` and subtracting it from every ``intercept`` leaves every prediction + unchanged -- so ``lstsq`` returns the minimum-norm representative and the + caller pins the constant afterwards. + + Alternating medians were tried here and are wrong: the grid is unbalanced + (each token count is reached by a different subset of batch sizes), so the + per-column median correction differs from column to column and the iteration + settles on a curve that is *not* the additive truth even when the data is + exactly additive. Robustness is recovered by trimming outliers and refitting + instead; see :func:`fit_additive_cost_model`. + """ + batch_index = {bs: i for i, bs in enumerate(sorted({c.batch_size for c in cells}))} + token_index = {m: i for i, m in enumerate(sorted({c.total_verify_tokens for c in cells}))} + num_batches = len(batch_index) + design = np.zeros((len(cells), num_batches + len(token_index)), dtype=np.float64) + observed = np.empty(len(cells), dtype=np.float64) + for row, cell in enumerate(cells): + design[row, batch_index[cell.batch_size]] = 1.0 + design[row, num_batches + token_index[cell.total_verify_tokens]] = 1.0 + observed[row] = cell.step_time_ms + + solution, *_ = np.linalg.lstsq(design, observed, rcond=None) + residuals = observed - design @ solution + intercept = {bs: float(solution[i]) for bs, i in batch_index.items()} + theta = {m: float(solution[num_batches + i]) for m, i in token_index.items()} + return intercept, theta, residuals + + +def fit_additive_cost_model( + cells: Sequence[CellStat], + *, + additive_tolerance: float = DEFAULT_ADDITIVE_TOLERANCE, +) -> AdditiveFit: + """Separate ``bias + alpha(bs)`` from ``theta(M)``, and report the misfit. + + A badly measured cell is deliberately *not* repaired. The residual of the + additive solve is the model-validity check the study calls for: cells that + share an ``M`` but come from different ``(bs, L)`` -- e.g. ``(32, 3)`` and + ``(16, 7)`` -- must agree once their intercepts are removed, and when they + do not, ``theta`` is not a function of ``M`` alone. Silently dropping the + disagreeing cell would delete exactly that evidence and leave a table that + fits beautifully and predicts badly. So the misfit is surfaced as a warning + and the human decides whether to trust the decomposition. + + The returned split is normalized so that ``theta`` extrapolates to zero at + ``M = 0``; see the module docstring for why that is the split the planner + needs rather than an arbitrary one. + """ + if not cells: + raise SweepGeometryError("no measured cells to fit") + _assert_grid_is_connected(cells) + + warnings: List[str] = [] + intercept, theta, residuals = _solve_additive(cells) + # Kept from the raw solve: this is the additive model's own error, and the + # shaping below (monotone clamp, floor clamp) is a deliberate deviation from + # the data that has its own warnings. Folding the two together would let a + # clamp masquerade as a model failure and vice versa. + typical = statistics.median(cell.step_time_ms for cell in cells) + max_abs = float(np.max(np.abs(residuals))) if len(residuals) else 0.0 + max_rel = max_abs / typical if typical > 0 else 0.0 + if max_rel > float(additive_tolerance): + worst = ", ".join( + f"(bs={cells[i].batch_size}, L={cells[i].verify_len}): {residuals[i]:+.3f} ms" + for i in np.argsort(-np.abs(residuals))[:4] + ) + warnings.append( + f"the additive model misses at least one cell by {max_abs:.3f} ms " + f"({max_rel:.1%} of a typical step), above the " + f"{float(additive_tolerance):.1%} tolerance. Largest disagreements: " + f"{worst}. A least-squares solve spreads one bad measurement over " + f"every cell sharing its batch size or token count, so read these as " + f"a neighbourhood rather than as the culprit. Either one of those " + f"steps did not replay a CUDA graph, or theta is not a function of M " + f"alone -- in which case the decomposed encoding is unsafe away from " + f"the profiled batch sizes." + ) + + ordered_tokens = sorted(theta) + monotone = running_max([theta[m] for m in ordered_tokens]) + if any(abs(monotone[i] - theta[m]) > 1e-9 for i, m in enumerate(ordered_tokens)): + warnings.append( + "theta(M) was not monotonically increasing and has been clamped " + "upward; a dip means the measurement noise is comparable to the " + "risers, so treat the finer shelves as unreliable." + ) + theta = dict(zip(ordered_tokens, monotone)) + + # Fix the one free constant: push theta's M -> 0 extrapolation onto the + # intercept, so what stays behind in `intercept` is the M-independent floor. + shift = _least_squares_intercept(ordered_tokens, [theta[m] for m in ordered_tokens]) + # ...but never so far that a shelf goes non-positive; SpsCostTable rejects + # that, and a zero-cost shelf would read as "these tokens are free". + shift = min(shift, min(theta.values()) - MIN_THETA_MS) + theta = {m: v - shift for m, v in theta.items()} + intercept = {bs: v + shift for bs, v in intercept.items()} + + if min(intercept.values()) < 0.0: + warnings.append( + "the fitted M-independent floor came out negative at the smallest " + "profiled batch size and has been clamped to zero; theta is then " + "explaining more than the whole step, which means the additive " + "model does not hold over this sweep." + ) + intercept = {bs: max(v, 0.0) for bs, v in intercept.items()} + + return AdditiveFit( + intercept_ms=dict(sorted(intercept.items())), + theta_ms=dict(sorted(theta.items())), + max_abs_residual_ms=max_abs, + max_rel_residual=max_rel, + warnings=tuple(warnings), + ) + + +def compress_to_risers( + token_counts: Sequence[int], + step_times: Sequence[float], + *, + min_riser_ms: float, + max_breakpoints: int, +) -> Tuple[List[int], List[float]]: + """Collapse a sampled curve into shelf-and-riser breakpoints. + + ``SpsCostTable`` interpolates between breakpoints, so what a table CLAIMS + is entirely in which points it keeps: dropping a shelf's interior is fine + (interpolation across equal values is still flat), but dropping a shelf's + LAST point turns the measured shelf into one long rising segment and + over-bills every mid-shelf total by up to the riser height -- the + over-trim mirror of the floor-lookup bug. So each kept riser is preceded + by the last point of the shelf it rises from: shelves survive as + breakpoint pairs, exactly the encoding the planner tests use for + "genuinely flat". + + A riser is kept only when it clears the previous kept shelf by + ``min_riser_ms``. The threshold is **absolute**, not relative to ``theta``: + ``theta`` is normalized to approach zero at ``M = 0``, so at the low end a + fraction of a millisecond of noise is a large *relative* move and a relative + threshold would manufacture shelves exactly where the curve carries least + information. What matters is whether a riser is worth a measurable fraction + of a whole step, which is what the caller scales this by. + + If more points survive than ``max_breakpoints``, the shallowest risers are + dropped first -- they are the ones a wrong call costs least. + """ + if len(token_counts) != len(step_times): + raise ValueError("token_counts and step_times must have the same length") + if not token_counts: + raise ValueError("nothing to compress") + + kept = [0] + for index in range(1, len(token_counts)): + if step_times[index] - step_times[kept[-1]] > float(min_riser_ms): + # Close the shelf we are leaving before recording the riser, so + # interpolation stays flat across it instead of ramping from the + # shelf's first sample to the riser top. + if index - 1 != kept[-1]: + kept.append(index - 1) + kept.append(index) + + limit = max(int(max_breakpoints), 2) + truncated_from = len(kept) + while len(kept) > limit: + # Drop the interior point whose removal introduces the least + # interpolation error. The previous rule -- smallest incoming jump -- + # deleted shelf-CLOSING points first by construction (their jump is + # below min_riser_ms), which turned every measured shelf back into a + # ramp and re-billed mid-shelf totals upward: the exact over-charge + # the pair encoding exists to prevent. Endpoints are never droppable: + # index 0 anchors the below-range clamp, the last point the above. + def _removal_error(pos: int) -> float: + x0, x1, x2 = (token_counts[kept[pos - 1]], token_counts[kept[pos]], + token_counts[kept[pos + 1]]) + y0, y2 = step_times[kept[pos - 1]], step_times[kept[pos + 1]] + interp = y0 + (y2 - y0) * (x1 - x0) / max(x2 - x0, 1) + return abs(step_times[kept[pos]] - interp) + + drop = min(range(1, len(kept) - 1), key=_removal_error) + kept.pop(drop) + if truncated_from > len(kept): + print(f"[dspark-sps] compress_to_risers: truncated {truncated_from} " + f"breakpoints to {len(kept)} (max_breakpoints={max_breakpoints});" + f" the kept set minimizes interpolation error, but raise the cap" + f" if the residuals matter", file=sys.stderr) + + return [int(token_counts[i]) for i in kept], [float(step_times[i]) for i in kept] + + +# --------------------------------------------------------------------------- +# table assembly, loading and refusals +# --------------------------------------------------------------------------- + + +def load_cost_table(payload: dict) -> SpsCostTable: + """Build the table exactly the way ``dspark.py::_build_verify_planner`` does. + + Kept byte-for-byte in step with the loader so this module can round-trip its + own output through the real validation instead of a copy of it. + """ + return SpsCostTable( + token_counts=[int(v) for v in payload["token_counts"]], + step_time_ms=[float(v) for v in payload["step_time_ms"]], + fixed_overhead_ms=float(payload.get("fixed_overhead_ms", 0.0)), + batch_sizes=[int(v) for v in payload.get("batch_sizes", [])], + batch_overhead_ms=[float(v) for v in payload.get("batch_overhead_ms", [])], + ) + + +def build_cost_table_payload( + cells: Sequence[CellStat], + *, + riser_tolerance: float = 0.01, + max_breakpoints: int = 8, + additive_tolerance: float = DEFAULT_ADDITIVE_TOLERANCE, + meta: Optional[dict] = None, +) -> dict: + """Assemble the JSON object ``confidence_sps_table_path`` expects. + + Emits exactly the five keys the loader reads plus a ``_meta`` block; unknown + keys are ignored by the loader, so provenance rides along for free. + + No leading ``token_counts = 0`` entry is synthesized. It would be dead + weight: :meth:`SpsCostTable.step_time` clamps a below-range query onto index + 0 already, so a fabricated base shelf could only either duplicate the first + measured shelf or extrapolate below everything that was measured. + """ + fit = fit_additive_cost_model(cells, additive_tolerance=additive_tolerance) + + # A riser has to be worth a measurable slice of a *whole step* to be real, + # so the noise floor is scaled by the typical step rather than by theta, + # whose own scale is an artifact of the M -> 0 normalization. + typical_step_ms = statistics.median(cell.step_time_ms for cell in cells) + min_riser_ms = float(riser_tolerance) * typical_step_ms + tokens, thetas = compress_to_risers( + sorted(fit.theta_ms), + [fit.theta_ms[m] for m in sorted(fit.theta_ms)], + min_riser_ms=min_riser_ms, + max_breakpoints=max_breakpoints, + ) + + batch_sizes = sorted(fit.intercept_ms) + intercepts = running_max([fit.intercept_ms[bs] for bs in batch_sizes]) + # Only the sum bias + alpha(bs) reaches the arithmetic, so the split is a + # reporting choice: bias is the floor that is present even at the smallest + # profiled batch, alpha is what each larger batch adds on top of it. + bias = min(intercepts) if intercepts else 0.0 + alphas = [value - bias for value in intercepts] + + payload = { + "token_counts": [int(v) for v in tokens], + "step_time_ms": [float(round(v, 6)) for v in thetas], + "fixed_overhead_ms": float(round(bias, 6)), + "batch_sizes": [int(v) for v in batch_sizes], + "batch_overhead_ms": [float(round(v, 6)) for v in alphas], + } + payload["_meta"] = { + # Consumer contract marker. Tables written for the old floor consumer + # deliberately dropped shelf-closing breakpoints; interpolating those + # re-bills every mid-shelf total upward with nothing to say so. The + # loader warns when this key is absent. + "lookup": "interp", + "encoding": "decomposed", + "note": ( + "step_time_ms is theta(M) only -- the trimmable term. The " + "non-trimmable floor is fixed_overhead_ms + batch_overhead_ms[bs]. " + "Do not read step_time_ms as whole-step time." + ), + "token_unit": ( + "total_verify_tokens = num_requests * (verify_len + 1); the bonus token is included" + ), + "cells": [ + cell.to_json() for cell in sorted(cells, key=lambda c: (c.batch_size, c.verify_len)) + ], + "fit": { + "max_abs_residual_ms": round(fit.max_abs_residual_ms, 6), + "max_rel_residual": round(fit.max_rel_residual, 6), + "additive_tolerance": float(additive_tolerance), + "riser_tolerance": float(riser_tolerance), + "min_riser_ms": round(min_riser_ms, 6), + "typical_step_ms": round(typical_step_ms, 6), + "warnings": list(fit.warnings), + }, + **(meta or {}), + } + # Constructing it here means a malformed table fails in the profiler rather + # than at engine init on a machine that has already spent minutes loading + # weights. + load_cost_table(payload) + return payload + + +def profitability_probe( + payload: dict, + *, + batch_size: int, + tiers: Sequence[int], + acceptance_rates: Sequence[float], +) -> List[dict]: + """Which verify length the real planner picks, per synthetic acceptance rate. + + Runs :func:`~.dspark_planner.budget_argmax_over_uniform_lens` -- the actual + production argmax, not a re-derivation -- against a geometric survival curve + ``p**(j+1)``. This is the question the table exists to answer and the one + a single sweep cannot: at what acceptance rate does + trimming become the better trade, and does it ever? + """ + table = load_cost_table(payload) + ladder = sorted({int(t) for t in tiers}) + block = max(ladder) + results: List[dict] = [] + for rate in acceptance_rates: + survival = np.tile( + np.array([float(rate) ** (j + 1) for j in range(block)], dtype=np.float64), + (int(batch_size), 1), + ) + chosen = budget_argmax_over_uniform_lens( + survival=survival, + num_gen_requests=int(batch_size), + cost_table=table, + allowed_lens=ladder, + min_verify_len=1, + ) + results.append( + { + "batch_size": int(batch_size), + "acceptance_rate": float(rate), + "chosen_verify_len": int(chosen), + "max_tier": int(ladder[-1]), + "trims": bool(int(chosen) < ladder[-1]), + } + ) + return results + + +def check_table_is_informative( + payload: dict, + *, + batch_sizes: Sequence[int], + tiers: Sequence[int], + acceptance_rates: Sequence[float] = (0.3, 0.5, 0.7, 0.8, 0.9, 0.95), + min_step_time_spread: float = 0.02, + allow_inert: bool = False, +) -> dict: + """Refuse to ship a table that cannot change the planner's mind. + + Two distinct failures, both of which produce a healthy-looking run: + + * *flat* -- ``is_flat`` is exact float equality, so it only catches a table + nobody measured. A measured-but-motionless curve passes it and then makes + ``decide_verify_lens`` return ``None`` on every step anyway. So the spread + of the **total** step time across the reachable token range is checked + too, not the spread of theta: theta is normalized towards zero at + ``M = 0``, which would make its own relative spread meaninglessly large. + * *inert* -- the curve moves, but every selectable tier at every profiled + batch size lands on one shelf. ``Theta = tau / T`` is then strictly + increasing in length and the argmax is always ``max_tier``. Unlike the + flat case this does *not* increment ``fallback_flat_cost``, so nothing + downstream would report it. + """ + table = load_cost_table(payload) + ladder = sorted({int(t) for t in tiers}) + if len(ladder) < 2: + raise InertCostTableError( + f"the verify-length ladder {ladder} has a single entry, so the " + f"planner has nothing to choose between and no cost table can " + f"change that." + ) + + if table.is_flat: + raise FlatCostTableError( + f"the profiled step_time_ms {payload['step_time_ms']} has a single " + f"distinct value, so SpsCostTable.is_flat is True. The planner " + f"treats that as 'unprofiled': decide_verify_lens returns None on " + f"every step (fallback_flat_cost) and ragged verification is a " + f"no-op. Either the sweep never crossed a cost riser -- widen it " + f"with larger batch sizes, since gains only exist at high " + f"concurrency -- or the step cost genuinely does not depend on the " + f"verified token count here, in which case confidence scheduling " + f"has no headroom on this configuration and should stay off." + ) + + spreads = {} + for batch_size in sorted({int(b) for b in batch_sizes}): + low = table.step_time(total_verify_tokens(batch_size, ladder[0]), batch_size) + high = table.step_time(total_verify_tokens(batch_size, ladder[-1]), batch_size) + spreads[batch_size] = (high - low) / low if low > 0 else 0.0 + best_spread = max(spreads.values(), default=0.0) + if best_spread < float(min_step_time_spread): + rendered = ", ".join(f"bs={k}: {v:.3%}" for k, v in sorted(spreads.items())) + raise FlatCostTableError( + f"the total step time only varies by {best_spread:.3%} across the " + f"whole tier ladder {ladder} at the best profiled batch size " + f"({rendered}), which is at or " + f"below the measurement noise floor of " + f"{float(min_step_time_spread):.3%}. Such a table is technically " + f"non-flat but the risers it asserts are noise, and the planner " + f"would trim on them. To be worth trimming, the step must grow by " + f"more than the acceptance yield given up -- typically >15% between " + f"adjacent tiers." + ) + + probes: List[dict] = [] + for batch_size in sorted({int(b) for b in batch_sizes}): + probes.extend( + profitability_probe( + payload, + batch_size=batch_size, + tiers=ladder, + acceptance_rates=acceptance_rates, + ) + ) + trims_anywhere = any(probe["trims"] for probe in probes) + if not trims_anywhere and not allow_inert: + raise InertCostTableError( + f"the table is non-flat but inert: over batch sizes " + f"{sorted({int(b) for b in batch_sizes})}, tiers {ladder} and " + f"acceptance rates {list(acceptance_rates)}, the planner picks " + f"max_tier={ladder[-1]} every time, so it would never trim a single " + f"token. That happens when no cost riser falls strictly inside " + f"(bs*(t_i+1), bs*(t_i+1 +1)] for any adjacent tier pair -- the " + f"tiers all sit on one shelf. Fix the ladder so its steps straddle " + f"the measured risers (see _meta." + f"recommended_confidence_verify_len_tiers), or accept that trimming " + f"cannot pay here. Pass --allow-inert to emit anyway." + ) + + return { + "step_time_spread_by_batch_size": {k: round(v, 6) for k, v in spreads.items()}, + "best_step_time_spread": round(best_spread, 6), + "profitability_probe": probes, + "trims_somewhere": trims_anywhere, + } + + +# --------------------------------------------------------------------------- +# the sweep +# --------------------------------------------------------------------------- + + +@dataclass +class SweepConfig: + """Everything the measurement loop needs; mirrors the CLI one-to-one.""" + + model: str + speculative_model: str + batch_sizes: List[int] + verify_lens: List[int] + tp_size: int = 1 + ep_size: Optional[int] = None + enable_attention_dp: bool = False + max_draft_len: int = 5 + input_len: int = 1024 + # 128, not a token gesture: autotuner cache misses ("using the fallback + # tactic") and allocator growth are still visible well past step 16 on a + # fresh engine, and a cell that starts on the cool side biases theta UP + # for exactly the small-M cells the argmax is most sensitive to. + warmup_steps: int = 128 + measure_steps: int = 64 + repeats: int = 1 + min_samples: int = 16 + max_seq_len: int = 4096 + # Frac sweeps only: override the engine's DERIVED tier ladder (cost-table + # shelf edges, typically [1,3,5]). A denser ladder makes the frac knob + # quantise to more rungs, so one sweep measures more M columns per bs -- + # the dense-tier experiments need the odd ratios {3,5} the default ladder + # cannot execute. + engine_tiers: Optional[List[int]] = None + max_num_tokens: int = 8192 + kv_cache_fraction: float = 0.5 + #: Fitted STS temperatures. Not used to decide anything while the pin is on; + #: it is what lets ``enable_ragged_verify`` pass validation without the cost + #: table this tool exists to produce. + sts_path: Optional[str] = None + #: Any non-flat cost table. See the note at its use site: with a one-rung + #: ladder its values cannot affect the result, only its presence. + seed_table_path: Optional[str] = None + #: Restore the original sweep geometry; see the note in _build_llm. + block_follows_verify_len: bool = False + #: Frac sweep (run_frac_sweep): budget fractions to visit on the live + #: ragged path. When set, verify_lens is pinned to [max_draft_len] by the + #: CLI and the per-cell M comes from measurement instead of the pin. + fracs: Optional[List[float]] = None + attn_backend: str = "TRTLLM" + disable_overlap_scheduler: bool = False + #: On by default because the deployment being profiled runs with it on, and + #: the confidence head is real per-step work that belongs in the table. + enable_confidence_scheduling: bool = True + timing_key: str = HOST_STEP_TIME_KEY + pin_acceptance: float = BONUS_TOKEN_ONLY_ACCEPTANCE + #: How long the post-cell stats drain waits for the tail records to arrive + #: over IPC before it decides the queue is empty. + stats_timeout_s: float = 5.0 + #: Mid-cell poll interval. Short enough that ~1000 rows cannot accumulate + #: between polls (at dp_size=8 that is 125 iterations, i.e. a few seconds). + stats_poll_s: float = 2.0 + token_id_range: Tuple[int, int] = (1000, 10000) + seed: int = 1234 + extra_llm_api_options: Dict[str, object] = field(default_factory=dict) + + @property + def dp_size(self) -> int: + """Independent scheduling domains, i.e. rows per iteration in the stats.""" + return int(self.tp_size) if self.enable_attention_dp else 1 + + def max_tokens_for(self, verify_len: int, + batch_size: Optional[int] = None) -> int: + """Output-token budget that guarantees the cell reaches its step count. + + With acceptance pinned every step commits exactly one token, so the + budget *is* the step count. Unpinned, a step commits up to + ``verify_len + 1``, so the budget is scaled to keep the lower bound on + steps at the target -- at the price of running up to that factor longer + when acceptance turns out to be poor. + + ``batch_size`` adds the *admission ramp*, and without it no large cell + can ever yield a sample. A cell is only counted while all ``dp_size`` + ranks hold exactly ``batch_size`` generating requests, so the request + admitted first must still be alive when the last one is admitted. + Measured at bs=256, input_len=1024: the scheduler admits + roughly one context request per iteration per rank, so the ramp is + ~``batch_size`` iterations, while the budget was 85 -- the first + requests died ~170 iterations before the last arrived, concurrency + plateaued at 11 of 256, and every one of the 8 repeats kept nothing. + bs<=128 escaped only because its ramp fit inside the same 85. + + The ramp is bounded below by prefill tokens / ``max_num_tokens``, but + that bound was off by ~60x here, so use the observed one-per-iteration + rate instead: it is the conservative direction, and overshooting costs + run time while undershooting costs the entire cell. + """ + target_steps = int(self.warmup_steps) + int(self.measure_steps) + budget = target_steps + max(target_steps // 4, 16) + if batch_size is not None: + budget += int(batch_size) + return budget if self.pin_acceptance else budget * (int(verify_len) + 1) + + def validate(self) -> None: + """Catch the geometry mistakes that would otherwise fail mid-sweep.""" + if self.fracs is not None: + bad = [f for f in self.fracs if not 0.0 < float(f) <= 1.0] + if bad: + raise SweepGeometryError( + f"--fracs values must be in (0, 1], got {sorted(bad)}") + if self.block_follows_verify_len: + raise SweepGeometryError( + "--fracs measures the live ragged path at the constant " + "block; --block-follows-verify-len contradicts it") + if not self.batch_sizes or not self.verify_lens: + raise SweepGeometryError("both --batch-sizes and --verify-lens must be non-empty") + if min(self.verify_lens) < 1: + raise SweepGeometryError("verify lengths must be >= 1 (the planner's floor)") + # A request that hits max_seq_len is evicted mid-cell, so the batch + # drains and the steady window shrinks -- which surfaces much later as a + # confusing "not enough steady samples". + longest = self.input_len + max( + self.max_tokens_for(v, b) for v in self.verify_lens for b in self.batch_sizes) + if longest > self.max_seq_len: + raise SweepGeometryError( + f"a request would reach {longest} tokens (input {self.input_len} + " + f"generated) but --max-seq-len is {self.max_seq_len}; raise it or " + f"lower --input-len / --measure-steps" + ) + # get_stats is served from a queue that retains on the order of 1000 + # ROWS, and attention-DP publishes one row per rank per iteration. A + # cell longer than that window keeps only its tail, which is the drain + # phase -- every step there has ranks disagreeing on numGenRequests, so + # the cell yields nothing and looks like it never ran. Raising + # --measure-steps past the window therefore makes things strictly worse; + # use --repeats to get more samples instead. + rows_per_iteration = self.dp_size + window_iterations = _STATS_ROW_BUDGET // max(rows_per_iteration, 1) + cell_iterations = int(self.warmup_steps) + int(self.measure_steps) + if cell_iterations > window_iterations: + print( + f"[dspark-sps] WARNING: a cell runs {cell_iterations} iterations " + f"but only about {window_iterations} fit in the retained stats " + f"window ({_STATS_ROW_BUDGET} rows / {rows_per_iteration} ranks). " + f"Only the tail survives, and the tail is the drain phase, so " + f"cells will come back empty. Lower --measure-steps to " + f"<= {window_iterations - int(self.warmup_steps)} and raise " + f"--repeats instead.", + file=sys.stderr, + ) + if self.input_len > self.max_num_tokens: + raise SweepGeometryError( + f"--input-len {self.input_len} exceeds --max-num-tokens " + f"{self.max_num_tokens}, so a prompt cannot be scheduled" + ) + widest_decode = max(self.batch_sizes) * (max(self.verify_lens) + 1) + if widest_decode > self.max_num_tokens: + raise SweepGeometryError( + f"the widest decode step submits {widest_decode} tokens " + f"(bs={max(self.batch_sizes)} x verify_len={max(self.verify_lens)} + 1) " + f"but --max-num-tokens is {self.max_num_tokens}; the batch would be " + f"split and no step would ever have the shape it is filed under" + ) + + +#: Approximate number of stats rows ``get_stats`` retains for one drain. Not a +#: constant the runtime exports; measured as 1001 rows surviving a 1300-step +#: attention-DP cell, i.e. the last 125 of its iterations. +_STATS_ROW_BUDGET = 1000 + + +def _prepare_environment(config: SweepConfig) -> None: + """Pin the knobs the table's meaning depends on, loudly. + + Set before the LLM is constructed because the worker processes inherit the + environment at spawn and read these once. + """ + # COMPACT, not STATIC. Under STATIC the token axis is bs*(max_draft_len+1) + # and the only way to vary M is to rebuild with a shorter block -- which + # also shortens the DRAFT pass, so the measured difference between two + # cells mixes "verified fewer tokens" with "drafted fewer tokens". The + # deployed path drafts the full block every step and varies only the verify + # window, so a table built that way over-states what trimming buys. COMPACT + # plus the per-engine pin below holds the draft at max_draft_len and moves + # only the axis the planner actually controls. + os.environ[RAGGED_VERIFY_MODE_ENV] = RaggedVerifyMode.COMPACT.value + if config.pin_acceptance: + from tensorrt_llm._torch.speculative.interface import ( + FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR) + + os.environ[FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR] = repr(float(config.pin_acceptance)) + + +def _build_llm(config: SweepConfig, verify_len: int, *, + pin_verify_len: bool = True): + """One engine per verify length; see the module docstring for why.""" + from tensorrt_llm import LLM + from tensorrt_llm.llmapi import CudaGraphConfig, DSparkDecodingConfig, KvCacheConfig + + # Held at the checkpoint's trained block for EVERY cell. Sweeping it (which + # this did) makes the draft pass shrink along with the verify window, and + # the two costs are then inseparable in the fitted theta(M). The previous + # K-sweep on this hardware put the draft pass at ~17 ms and the marginal + # verify token at ~0 -- so a table that lets draft cost leak into theta(M) + # would attribute the draft saving to trimming and predict a benefit the + # deployment cannot collect. The verify window is moved by the pin below. + # The window is pinned by BOTH the env override and the ladder, and the + # env var is SET here, never popped. + # + # Set, because the workers are spawned by srun with the job environment: a + # job-wide pin reaches ranks 1..N-1 regardless of what this process does to + # its own env, so popping it here strips rank 0 alone. A rank-0 planner + # with no pin and no cost table falls back to uniform while its peers run + # pinned -- a rank-divergent shape under attention-DP. With one verify + # length per sweep process the pin always equals the ladder rung below, so + # it can never be rejected either. + from tensorrt_llm._torch.speculative.dspark_verify import FORCE_VERIFY_LEN_ENV + + # LAUNCH-MODE CAVEAT: this mutation reaches ranks 1..N-1 only when the + # workers are spawned AFTER it (MPI-spawn path). On a pre-spawned world + # (srun / trtllm-llmapi-launch) already-running ranks keep the previous + # value, the per-L pin diverges across ranks, and cells measure a shape + # their label does not describe. assert_swept_shape refuses such a sweep + # downstream (verify_len_hist vs the nominal L), which is why it gates the + # fit rather than merely warning. + # Frac sweeps leave the pin unset: their windows come from the live + # planner (that is the experiment), and the budget fraction travels by + # RPC + the step's own allgather, which reaches pre-spawned worlds the + # env mutation below cannot. + if pin_verify_len: + os.environ[FORCE_VERIFY_LEN_ENV] = str(int(verify_len)) + else: + os.environ.pop(FORCE_VERIFY_LEN_ENV, None) + # Which knob moves M. + # + # constant block (default): the block stays at the checkpoint's trained + # length for every cell and only the verify window moves, which is exactly + # what the deployment does. + # + # block follows verify_len: the original design. Never produces + # max(lens) < block, so it runs today, at the price of shrinking the DRAFT + # pass along with the window -- theta(M) then absorbs a per-L draft term. + # That term is recoverable rather than fatal: with several batch sizes the + # same M is reached at different L (M=192 from bs=32/L=5 and bs=64/L=2), and + # those shared points separate a token cost from an L cost. + block = int(verify_len) if config.block_follows_verify_len else int(config.max_draft_len) + spec_config = DSparkDecodingConfig( + max_draft_len=block, + block_size=block, + speculative_model=config.speculative_model, + enable_confidence_scheduling=bool(config.enable_confidence_scheduling), + # Both required together: the validator rejects confidence scheduling + # without ragged verification (there are only two states -- schedule per + # request, or verify the full block). The profiler had only the first + # and could not construct its own engine at all; it failed 1:36 into + # every two-node sweep. And ragged in turn demands a profiled artifact, + # which is circular for the tool that produces one -- an STS path is the + # non-circular way to satisfy it, and the pin makes the planner's own + # choice irrelevant here anyway. + enable_ragged_verify=bool(config.enable_confidence_scheduling), + confidence_sts_path=config.sts_path, + # Only has to be non-flat and present: with one rung the planner's + # choice is already determined, so the table's VALUES cannot influence + # this measurement -- it exists to get past the flat-cost fallback, + # which would otherwise return None and put every step back on the + # uniform path at M = bs*(block+1). + confidence_sps_table_path=config.seed_table_path, + # A single tier keeps the captured-graph count identical to an + # unscheduled run (one graph per batch size) while still exercising the + # confidence head, which is part of the deployment's per-step cost. It + # also makes the planner's choice a no-op: max_tier == max_draft_len, so + # the step shape is pinned no matter what the confidence snapshot says. + # The pin decides the length; the ladder only has to CONTAIN it, or the + # planner would return a length with no captured graph and every step + # would fall back to eager -- which is not the cost being tabulated. + # One rung: the planner's argmax has a single legal answer, so the step + # shape is pinned no matter what the confidence snapshot says, while + # every downstream path (bucket choice, padding reservation, graph + # selection) runs exactly as it does in production. + confidence_verify_len_tiers=( + # Frac sweeps need the PRODUCTION ladder: the token-bucket grid is + # derived from it, and a one-rung ladder would collapse every + # captured bucket to the full block -- the sweep would then only + # ever measure M = bs * (block + 1) regardless of the frac. + # ``engine_tiers`` overrides the derived ladder for dense grids. + (([int(t) for t in config.engine_tiers] + if config.engine_tiers else None) + if not pin_verify_len else [int(verify_len)]) + if config.enable_confidence_scheduling else None + ), + ) + llm_kwargs = dict( + model=config.model, + tensor_parallel_size=int(config.tp_size), + moe_expert_parallel_size=int(config.ep_size) if config.ep_size else None, + enable_attention_dp=bool(config.enable_attention_dp), + attn_backend=config.attn_backend, + max_batch_size=max(config.batch_sizes), + max_seq_len=int(config.max_seq_len), + max_num_tokens=int(config.max_num_tokens), + disable_overlap_scheduler=bool(config.disable_overlap_scheduler), + speculative_config=spec_config, + # Block reuse would let a later cell start from a warm prefix and skip + # prefill work the earlier cells paid for, which shifts the steady-state + # KV length between cells. + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + free_gpu_memory_fraction=float(config.kv_cache_fraction), + ), + # Any batch size outside this list runs eager; an eager step is orders + # of magnitude off the graph-replay cost and would poison the curve. + cuda_graph_config=CudaGraphConfig( + batch_sizes=sorted(set(config.batch_sizes)), + enable_padding=True, + ), + enable_iter_perf_stats=True, + # Unbounded: the buffer is drained while the sweep runs, but a stall + # during a long cell must not silently evict the steps being measured. + iter_stats_max_iterations=-1, + ) + llm_kwargs = {k: v for k, v in llm_kwargs.items() if v is not None} + llm_kwargs.update(config.extra_llm_api_options) + return LLM(**llm_kwargs) + + +def _run_cell(llm, config: SweepConfig, *, batch_size: int, verify_len: int, + measure_ragged: bool = False) -> List[StepSample]: + """Hold ``batch_size`` generating requests per rank and record every step. + + With ``measure_ragged`` the per-step token total is READ from the stats + (the group's agreed bucket) instead of derived from the pin, and samples + carry it in ``measured_tokens``; one cell then legitimately yields samples + at several M values. + """ + from tensorrt_llm import SamplingParams + + rng = random.Random(config.seed + batch_size * 1000 + verify_len) + low, high = config.token_id_range + num_requests = int(batch_size) * config.dp_size + sampling_params = SamplingParams( + # Every request must outlive the whole window so the batch never drains + # mid-measurement and shrinks the shape being profiled. + max_tokens=config.max_tokens_for(verify_len, batch_size), + ignore_eos=True, + # No detokenization: it is postprocessing work that would show up in the + # host step time without being part of what the planner is choosing. + detokenize=False, + temperature=0.0, + ) + # Random ids from a mid-vocabulary window rather than real text: no tokenizer + # round-trip is needed, and a distinct prompt per request keeps any prefix + # sharing out of the measurement. The window avoids the low ids where + # special/control tokens live in every large vocabulary. + prompts = [ + {"prompt_token_ids": [rng.randint(low, high) for _ in range(config.input_len)]} + for _ in range(num_requests) + ] + + futures = [llm.generate_async(prompt, sampling_params) for prompt in prompts] + + # Drained WHILE the cell runs, not once at the end. Only ~1000 stats rows + # survive to a single drain -- 125 iterations at dp_size=8 -- and the + # scheduler admits roughly one context request per iteration, so a cell + # spends ~batch_size iterations ramping up and just as many draining. At + # bs>=128 the whole retained window therefore lands inside the drain: every + # step it contains has 0..11 generating requests, none has `batch_size`, + # and the cell reports "kept nothing" while the plateau it was measuring + # scrolled out of the buffer unseen. + # + # `IterationResult` latches done when its queue runs dry, so an + # empty read ends collection. Polling only while requests are still + # outstanding keeps the queue fed; the count of empty reads is reported so a + # cell that lost its tail this way is visible rather than merely short. + rows: List[dict] = [] + drain_error: List[BaseException] = [] + empty_reads = [0] + stop = threading.Event() + + def _drain() -> None: + while not stop.is_set(): + try: + got = list(llm.get_stats(timeout=config.stats_poll_s) or []) + except BaseException as exc: # noqa: BLE001 - reported, not swallowed + drain_error.append(exc) + return + if got: + rows.extend(got) + else: + # Counted because one empty read is not merely a lost poll: + # IterationResult latches done on the first Empty, and until + # something un-latches it every later get_stats -- including + # the final post-cell drain -- returns [] instantly. A cell + # with a non-zero count here may have lost its tail. + empty_reads[0] += 1 + stop.wait(config.stats_poll_s) + + # A daemon drainer, NOT a poll loop that owns the exit condition. Making the + # main thread wait on `future.finished` while draining hung the sweep: the + # loop span silently, the log stopped growing, and the stall watchdog killed + # the job at exactly 15 minutes with no cell produced. The main thread here + # blocks on future.result() exactly as it did before this change, so + # termination is unaffected by anything the drainer does. + # stats_poll_s <= 0 disables the mid-cell drain: the control arm for "does + # the profiler perturb what it measures". The cell then relies on the single + # post-cell drain, which only sees the last ~1000 rows -- fine for small + # batches whose whole cell fits in that window, useless above bs~64. That is + # exactly why this is a control and not an option. + drainer = None + if float(config.stats_poll_s) > 0: + drainer = threading.Thread(target=_drain, name="dspark-sps-drain", + daemon=True) + drainer.start() + try: + for future in futures: + future.result() + finally: + stop.set() + if drainer is not None: + drainer.join(timeout=max(5.0, 3.0 * config.stats_poll_s)) + for future in futures: + if not future.finished: + future.abort() + if drain_error: + print(f"[dspark-sps] bs={batch_size} L={verify_len}: mid-cell drain " + f"stopped early ({type(drain_error[0]).__name__}: " + f"{drain_error[0]}); the cell may be short", + file=sys.stderr) + # Final drain for whatever crossed the IPC boundary after the last poll. + # Un-latch first: a single gap in stats production longer than the poll + # timeout latched IterationResult done, and a latched queue answers [] + # instantly -- silently truncating the cell this drain exists to finish. + undo = getattr(getattr(getattr(llm, "_executor", None), + "_iter_stats_result", None), "mark_undone", None) + if undo is not None: + undo() + rows.extend(list(llm.get_stats(timeout=config.stats_timeout_s) or [])) + if empty_reads[0]: + print(f"[dspark-sps] bs={batch_size} L={verify_len}: " + f"{empty_reads[0]} empty stats read(s) mid-cell; the retained " + f"window may have closed early and the cell lost its tail", + file=sys.stderr) + + if measure_ragged: + ragged = ragged_steps_from_stats( + rows, expected_ranks=config.dp_size, timing_key=config.timing_key) + kept_ragged = [s for s in ragged if s[1] == int(batch_size)] + if not kept_ragged: + print( + f"[dspark-sps] bs={batch_size:>5} (ragged) kept nothing -- " + + explain_alignment(rows, expected_ranks=config.dp_size, + timing_key=config.timing_key) + + "; also requires numGenerationTokens on every rank row", + file=sys.stderr) + return [] + # Warmup is discarded per CELL by iteration order, here rather than in + # summarize_cells: one ragged cell yields several (bs, M) groups, and a + # per-group trim would re-charge the warmup against every bucket. + cutoff = sorted(s[0] for s in kept_ragged) + cutoff = cutoff[min(int(config.warmup_steps), len(cutoff) - 1)] + return [ + StepSample( + batch_size=int(batch_size), + verify_len=int(verify_len), + step_time_ms=step_time, + iteration=iteration, + measured_tokens=tokens, + ) + for iteration, num_gen, tokens, step_time in kept_ragged + if iteration >= cutoff + ] + aligned = aligned_steps_from_stats( + rows, expected_ranks=config.dp_size, timing_key=config.timing_key + ) + kept = [s for s in aligned if s[1] == int(batch_size)] + if not kept: + # Printed per cell rather than once at the end: which cells failed and + # why is the whole diagnosis, and a sweep that ends in + # InsufficientSamplesError has already thrown that away. + print( + f"[dspark-sps] bs={batch_size:>5} L={verify_len} kept nothing -- " + + explain_alignment(rows, expected_ranks=config.dp_size, + timing_key=config.timing_key) + + (f"; aligned but numGenRequests never == {batch_size} " + f"(saw {sorted({a[1] for a in aligned})[:12]})" if aligned else ""), + file=sys.stderr, + ) + return [ + StepSample( + batch_size=int(batch_size), + verify_len=int(verify_len), + step_time_ms=step_time, + iteration=iteration, + ) + for iteration, num_gen, step_time in aligned + if num_gen == int(batch_size) + ] + + +def run_sweep(config: SweepConfig) -> List[StepSample]: + """Measure every ``(batch_size, verify_len)`` cell, rebuilding per length.""" + config.validate() + _prepare_environment(config) + samples: List[StepSample] = [] + for verify_len in sorted(config.verify_lens): + print( + f"[dspark-sps] building engine: block_size=" + f"{verify_len if config.block_follows_verify_len else config.max_draft_len}" + f"{' (follows L)' if config.block_follows_verify_len else ' (constant)'}" + f", verify window {verify_len} -> M = bs*{verify_len + 1} ...", + file=sys.stderr, + ) + llm = _build_llm(config, verify_len) + try: + for batch_size in sorted(config.batch_sizes): + started = time.time() + cell = [] + for _ in range(max(1, int(config.repeats))): + # Each repeat drains the stats queue again, so repeats add + # samples where a longer single cell would only overflow the + # retained window. + cell.extend(_run_cell(llm, config, batch_size=batch_size, + verify_len=verify_len)) + print( + f"[dspark-sps] bs={batch_size:>5} L={verify_len} " + f"M={total_verify_tokens(batch_size, verify_len):>6} " + f"aligned_steps={len(cell):>5} " + f"({time.time() - started:.1f}s)", + file=sys.stderr, + ) + samples.extend(cell) + finally: + llm.shutdown() + return samples + + +def run_frac_sweep(config: SweepConfig) -> List[StepSample]: + """Measure ``(batch_size, frac)`` cells on the LIVE ragged path. + + One engine build for the whole sweep: the block stays at the checkpoint's + trained length, no verify-length pin is set, and per cell the budget + fraction is switched at runtime (``set_dspark_budget_frac`` -- rank 0's + queue, delivered by the decode loop's own allgather). Windows come from + the real confidence top-k and the executed shape is the group's agreed + bucket, so the cells sample exactly the geometry production runs -- + including the token-major presentation and the bucket fit -- instead of + the pinned-uniform approximation. M is read per step from the stats + (``measured_tokens``), so one cell contributes samples at every bucket the + planner visited. + """ + config.validate() + _prepare_environment(config) + fracs = sorted({float(f) for f in (config.fracs or [])}) + if not fracs: + raise ValueError("run_frac_sweep requires a non-empty --fracs") + print( + f"[dspark-sps] building engine once: block_size={config.max_draft_len} " + f"(constant), live planner windows, frac sweep {fracs} ...", + file=sys.stderr, + ) + samples: List[StepSample] = [] + llm = _build_llm(config, int(config.max_draft_len), pin_verify_len=False) + + def _set_frac(value): + executor = llm._executor + rpc = getattr(executor, "collective_rpc", None) + if rpc is not None: + return rpc("set_dspark_budget_frac", args=(value,)) + # Single-process executor path: the worker IS the executor. + return executor.set_dspark_budget_frac(value) + + try: + for frac in fracs: + _set_frac(float(frac)) + for batch_size in sorted(config.batch_sizes): + started = time.time() + cell: List[StepSample] = [] + for _ in range(max(1, int(config.repeats))): + cell.extend( + _run_cell(llm, config, batch_size=batch_size, + verify_len=int(config.max_draft_len), + measure_ragged=True)) + buckets = sorted({s.measured_tokens for s in cell}) + print( + f"[dspark-sps] bs={batch_size:>5} frac={frac:.2f} " + f"aligned_steps={len(cell):>5} buckets={buckets[:8]}" + f"{' ...' if len(buckets) > 8 else ''} " + f"({time.time() - started:.1f}s)", + file=sys.stderr, + ) + samples.extend(cell) + finally: + try: + _set_frac(None) + except Exception: # noqa: BLE001 - shutdown path, engine may be gone + pass + llm.shutdown() + return samples + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def _int_list(text: str) -> List[int]: + return [int(v) for v in str(text).replace(",", " ").split()] + + +def _float_list(text: str) -> List[float]: + return [float(v) for v in str(text).replace(",", " ").split()] + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python tests/microbenchmarks/dspark_sps_profiler.py", + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--out", required=True, help="Path to write the cost-table JSON.") + source = parser.add_argument_group( + "sample source", + "Where the (batch size, verify length, step time) triples come from. " + "Default: build an engine in-process and sweep it. The two " + "alternatives measure the deployment itself, so the table cannot " + "describe a configuration nobody is running.") + source.add_argument( + "--from-iter-log", nargs="+", metavar="LOG", + help="Fit from production iteration logs (print_iter_log). No server, " + "no traffic, no machine time; coverage is whatever the " + "deployment happened to run.") + source.add_argument( + "--base-url", + help="Sweep a RUNNING trtllm-serve instead of building an engine. The " + "verify length is pinned over /dspark/verify_len_pin, so walking " + "the ladder costs no engine rebuild.") + source.add_argument( + "--control-url", + help="Where the pin endpoint and /metrics live when they differ from " + "--base-url: under disaggregated serving the LOAD goes through the " + "proxy, but /dspark/verify_len_pin and the iteration stats belong to " + "the GENERATION server. Defaults to --base-url.", + ) + source.add_argument("--server-settle-s", type=float, default=20.0, + help="Seconds to let admission reach the plateau " + "before a cell is measured (--base-url).") + source.add_argument("--server-poll-s", type=float, default=20.0, + help="Seconds of steady state to collect per cell " + "(--base-url).") + source.add_argument( + "--server-log", + help="The server's stdout/stderr (print_iter_log). When reachable this " + "is where --base-url takes its samples from: every step is " + "labelled by the shape it ran, and no metrics collector has to be " + "configured for the sweep to see anything.") + source.add_argument("--served-model-name", default=None, + help="`model` field for /v1/completions (--base-url). " + "Omit to read it from /v1/models, which is what " + "the server actually registered -- for a local " + "checkpoint that is the directory basename, not " + "the path it was launched with.") + source.add_argument( + "--cuda-graph-batch-sizes", + help="The deployment's captured batch-size ladder, comma separated " + "(--from-iter-log). Padding rounds a step's rows up to the " + "smallest entry at or above them, which is what makes the padded " + "token total resolve to exactly one verify length. Omit and only " + "steps whose shape is unambiguous on its own are kept.") + parser.add_argument( + "--samples-out", + default=None, + help="Raw per-step samples (JSONL). Defaults to .samples.jsonl.", + ) + parser.add_argument( + "--fit-only", + action="store_true", + help="Skip the engine entirely and refit from --samples-in. No GPU needed.", + ) + parser.add_argument( + "--samples-in", + default=None, + help="Comma-separated JSONL sample files to fit (implies --fit-only inputs).", + ) + + model = parser.add_argument_group("model") + model.add_argument("--model", help="Target model checkpoint path or HF id.") + model.add_argument( + "--speculative-model", + default=None, + help="DSpark draft checkpoint. Defaults to --model.", + ) + model.add_argument("--tp-size", type=int, default=1) + model.add_argument("--ep-size", type=int, default=None, help="moe_expert_parallel_size.") + model.add_argument("--enable-attention-dp", action="store_true") + model.add_argument("--attn-backend", default="TRTLLM") + model.add_argument("--max-seq-len", type=int, default=4096) + model.add_argument("--max-num-tokens", type=int, default=8192) + model.add_argument("--kv-cache-fraction", type=float, default=0.5) + model.add_argument( + "--disable-overlap-scheduler", + action="store_true", + help="Profile without the overlap scheduler. The default (overlap ON) is " + "what ships, so leave this off unless the deployment disables it.", + ) + model.add_argument( + "--no-confidence-scheduling", + dest="enable_confidence_scheduling", + action="store_false", + help="Profile with confidence scheduling off. Only for a checkpoint " + "without confidence-head weights: the head is real per-step work and " + "the deployment being profiled runs with it on.", + ) + model.add_argument( + "--extra-llm-api-options", + default=None, + help="YAML file merged into the LLM(...) kwargs (moe_config, quant, ...).", + ) + parser.set_defaults(enable_confidence_scheduling=True) + + sweep = parser.add_argument_group("sweep") + sweep.add_argument( + "--batch-sizes", + default="8,16,32,64,128", + help="Generating requests PER attention-DP rank. Every value is captured " + "as a CUDA graph; use a ladder whose ratios are reachable by the " + "verify-length ratios (a doubling ladder is safe).", + ) + sweep.add_argument( + "--max-draft-len", + type=int, + default=5, + help="The deployment's block size. Sets the top of the verify-length sweep.", + ) + sweep.add_argument( + "--verify-lens", + default=None, + help="Uniform verify lengths to sweep. Defaults to 1..--max-draft-len. " + "Each value costs one engine build.", + ) + sweep.add_argument( + "--fracs", + default=None, + help="Budget fractions in (0, 1] to sweep on the LIVE ragged path " + "(one engine build total; windows from the real confidence top-k, " + "per-step M measured from the group's agreed bucket). Mutually " + "exclusive with --verify-lens and --block-follows-verify-len.", + ) + sweep.add_argument( + "--engine-tiers", + default=None, + help="Frac sweeps only: comma-separated tier ladder for the SWEEP " + "engine (e.g. 2,3,4,5), overriding the cost-table-derived default. " + "A denser ladder quantises the frac knob to more rungs, so one sweep " + "measures more M columns per batch size.", + ) + sweep.add_argument("--input-len", type=int, default=1024, help="Prompt length in tokens.") + sweep.add_argument("--block-follows-verify-len", action="store_true", + help="Shrink the drafted block with the verify window " + "(the original geometry). Avoids the " + "max(lens) List[StepSample]: + samples: List[StepSample] = [] + for path in paths: + with open(path, encoding="utf-8") as handle: + for line in handle: + line = line.strip() + if line: + samples.append(StepSample.from_json(json.loads(line))) + return samples + + +def _write_samples(path: str, samples: Sequence[StepSample]) -> None: + with open(path, "w", encoding="utf-8") as handle: + for sample in samples: + handle.write(json.dumps(sample.to_json()) + "\n") + + +# --------------------------------------------------------------------------- +# Alternative sample sources +# +# The in-process sweep below builds its own engine per verify length. That +# gives total control, and costs a rebuild per length plus a configuration +# that has to be kept in sync with the deployment by hand -- when it drifts, +# the table describes a machine nobody is running (a profiled (256, 1536) cell +# read 150.2 ms against 118.8 ms measured live at the same shape, and nothing +# in the fingerprint recorded the difference). +# +# These two collectors take the measurement to the deployment instead. +# --------------------------------------------------------------------------- + + +def resolve_padded_shape( + *, + num_rows: int, + num_generation_tokens: int, + max_draft_len: int, + max_batch_size: int, + padded_batch_sizes: Optional[Sequence[int]] = None, +) -> Optional[Tuple[int, int]]: + """Recover ``(padded_bs, verify_len)`` from a step's row and token counts. + + The token total is the PADDED one: a step with 244 real rows submits + ``256 * (L + 1)`` tokens because the batch is padded up to a captured + graph size, and it is that padded shape the step actually cost. So the + length cannot be read off the real row count -- ``M / rows`` is not even + integral -- and a filter that demands it throws away nearly every step (it + kept 312 of 5682 on two production runs). + + Every candidate length implies a padded row count; the answer is the one + that is integral, at least the real row count, and within the engine's + batch ceiling. Ambiguous steps are rejected rather than guessed: filing a + step under a shape it did not run is the mislabelling that once produced a + clean-looking, entirely fictional table. + """ + if num_rows <= 0 or num_generation_tokens <= 0: + return None + ladder = (sorted({int(b) for b in padded_batch_sizes}) + if padded_batch_sizes else None) + if ladder: + # Padding rounds UP to the smallest captured batch size, never past + # it, so the row count alone fixes the padded width and the length + # follows. Without this the low-occupancy steps are genuinely + # ambiguous -- 768 tokens over 128 rows is L=5 at 128 rows, L=3 at + # 192, or L=2 at 256 -- and all of them would be dropped. + candidates = [b for b in ladder if b >= num_rows] + if not candidates: + return None + padded_bs = candidates[0] + if num_generation_tokens % padded_bs: + # Padding declined (a resource bail), so the step ran its real + # rows and this width does not describe it. Skip rather than + # invent one. + return None + verify_len = num_generation_tokens // padded_bs - 1 + if not 0 < verify_len <= int(max_draft_len): + return None + return (padded_bs, verify_len) + found: List[Tuple[int, int]] = [] + for verify_len in range(1, int(max_draft_len) + 1): + width = verify_len + 1 + if num_generation_tokens % width: + continue + padded_bs = num_generation_tokens // width + if padded_bs < num_rows or padded_bs > int(max_batch_size): + continue + if ladder is not None and padded_bs not in ladder: + continue + found.append((padded_bs, verify_len)) + return found[0] if len(found) == 1 else None + + +def samples_from_iter_log( + paths: Sequence[str], + *, + max_draft_len: int, + max_batch_size: int = 256, + padded_batch_sizes: Optional[Sequence[int]] = None, + rank_prefix: str = "0:", + max_step_ms: float = 1000.0, +) -> List[StepSample]: + """Read (bs, M, step time) triples out of production iteration logs. + + ``print_iter_log`` already emits everything a cell needs:: + + iter = N, ..., num_scheduled_requests = R, ..., + host_step_time = T ms, ..., states = {..., 'num_generation_tokens': M} + + so a table can be fitted from traffic that really ran, at the deployment's + own context depths, batch mix and captured-graph set -- no synthetic + prompts, no second engine, no machine time. The limitation is coverage: + only shapes the deployment happened to run appear, and rarely-taken rungs + may be thin. Use it to fit a table for a workload that is already in + production, or to check that a swept table has not drifted; use the server + sweep when specific cells have to be visited on purpose. + + ``verify_len`` is recovered by :func:`resolve_padded_shape`, which reads it + out of the padded token total instead of assuming it; steps whose shape is + ambiguous are skipped rather than filed under a length they did not run. + Samples carry the PADDED row count, because that is the shape whose cost + was measured. + + Pass ``padded_batch_sizes`` (the deployment's CUDA-graph ladder) when it is + known: it removes the last ambiguity, since only a captured row count can + have run. + """ + line_re = re.compile( + r"iter = (\d+).*?num_scheduled_requests = (\d+).*?" + r"host_step_time = ([\d.]+)ms.*?'num_ctx_requests': (\d+)" + r".*?'num_generation_tokens': (\d+)") + out: List[StepSample] = [] + skipped_ratio = 0 + skipped_mixed = 0 + + def _consume(line: str) -> None: + nonlocal skipped_ratio, skipped_mixed + match = line_re.search(line) + if not match: + return + iteration = int(match.group(1)) + batch_size = int(match.group(2)) + step_ms = float(match.group(3)) + ctx_requests = int(match.group(4)) + tokens = int(match.group(5)) + if ctx_requests: + # A step carrying prefill chunks runs eager at a prefill-inflated + # time; the ladder resolution would still find a plausible + # (padded_bs, L) for it and file a mixed step as a pure decode + # cell -- confidently mislabelled, the exact failure this + # collector exists to avoid. + skipped_mixed += 1 + return + if iteration == 1: + # Each executor instance numbers from 1, and its first step + # carries the wait for the first request inside its host time + # (62 s on a live server). Identify it by the number rather than + # by having seen one, so a log FRAGMENT is still usable. + return + if batch_size <= 0 or step_ms >= max_step_ms: + return + shape = resolve_padded_shape( + num_rows=batch_size, num_generation_tokens=tokens, + max_draft_len=max_draft_len, max_batch_size=max_batch_size, + padded_batch_sizes=padded_batch_sizes) + if shape is None: + skipped_ratio += 1 + return + padded_bs, verify_len = shape + out.append(StepSample(batch_size=padded_bs, verify_len=verify_len, + step_time_ms=step_ms, iteration=iteration)) + + for path in paths: + text = pathlib.Path(path).read_text(errors="replace") + candidates = [ln for ln in text.splitlines() if "iter = " in ln] + matched = [ln for ln in candidates + if not rank_prefix or ln.lstrip().startswith(rank_prefix)] + if candidates and not matched and rank_prefix: + # Single-process deployments write no srun/mpirun rank tag; a + # hard prefix requirement silently dropped every line and the + # error blamed print_iter_log. When the prefix matches nothing at + # all, the log simply is not rank-tagged -- read it as-is. + print(f"[dspark-sps] iter-log: no lines in {path} carry the rank " + f"prefix {rank_prefix!r}; reading it as an untagged " + f"single-process log", file=sys.stderr) + matched = candidates + for line in matched: + _consume(line) + if skipped_mixed: + print(f"[dspark-sps] iter-log: skipped {skipped_mixed} mixed " + f"context+generation steps (prefill-inflated times)", + file=sys.stderr) + if skipped_ratio: + print(f"[dspark-sps] iter-log: skipped {skipped_ratio} steps whose " + f"padded shape could not be resolved unambiguously", + file=sys.stderr) + return out + + +def samples_from_server( + *, + base_url: str, + control_url: Optional[str] = None, + batch_sizes: Sequence[int], + verify_lens: Sequence[int], + input_len: int, + model: str, + warmup_steps: int, + settle_s: float, + poll_s: float, + max_draft_len: int, + seed: int, + padded_batch_sizes: Optional[Sequence[int]] = None, + server_log: Optional[str] = None, + dp_size: int = 1, +) -> List[StepSample]: + """Sweep a RUNNING server: pin the window over HTTP, then read /metrics. + + The equivalent of SGLang's profiler loop, and the reason + ``/dspark/verify_len_pin`` exists: walking the verify-length ladder no + longer costs an engine rebuild per rung, and the engine being measured is + by construction the engine that will consume the table. + + Load is synthetic (random token ids, ignore_eos) because a cell has to hold + a fixed batch size for long enough to measure it; the SHAPE is what the + table indexes, and the deployment's own graphs, context depths and + scheduler are what answer for it. + + ``dp_size`` is how many requests it takes to put ONE on each rank's + batch. Under attention DP every rank schedules its own subset, the + planner prices its own rank's batch, and the iteration log reports + per-rank rows -- so the table's bs axis is PER-RANK, and a cell labeled + ``bs`` needs ``bs * dp_size`` requests in flight. Job 2585129 swept a + tp8-ADP server with dp_size unaccounted: every requested cell (64..256) + actually measured per-rank batches of 8..40, and the fit refused the + grid as disconnected. + """ + import requests + + control_url = control_url or base_url + rng = random.Random(seed) + out: List[StepSample] = [] + + if not model: + # The served id is whatever the server registered -- for a local + # checkpoint that is the directory's basename, not the path that was + # passed to it. Posting the path gets every request rejected, no + # traffic runs, and every cell reports zero steps with nothing to say + # why. Ask the server instead of making the caller guess. + listing = requests.get(f"{base_url}/v1/models", timeout=60) + listing.raise_for_status() + entries = (listing.json() or {}).get("data") or [] + if not entries: + raise RuntimeError(f"{base_url}/v1/models listed no model") + model = entries[0]["id"] + print(f"[dspark-sps] serving model id: {model}", file=sys.stderr) + + post_errors: List[str] = [] + + def _fire(body: dict) -> None: + try: + response = requests.post(f"{base_url}/v1/completions", json=body, + timeout=3600) + if response.status_code != 200 and len(post_errors) < 3: + post_errors.append( + f"{response.status_code} {response.text[:200]}") + except Exception as exc: # noqa: BLE001 - reported below, not fatal + if len(post_errors) < 3: + post_errors.append(f"{type(exc).__name__}: {exc}") + + def pin(verify_len: Optional[int]) -> None: + response = requests.post(f"{control_url}/dspark/verify_len_pin", + json={"verify_len": verify_len}, timeout=60) + if response.status_code != 200: + raise RuntimeError( + f"pinning verify_len={verify_len} was refused by the server: " + f"{response.status_code} {response.text}") + + def drain_metrics() -> List[dict]: + response = requests.get(f"{control_url}/metrics", timeout=60) + response.raise_for_status() + payload = response.json() + return payload if isinstance(payload, list) else [payload] + + try: + for verify_len in sorted(verify_lens): + try: + pin(int(verify_len)) + except RuntimeError as exc: + # A rung the server's ladder lacks. Raising out of the sweep + # threw away every already-collected sample (nothing is + # written until the sweep returns); a partial ladder is still + # a table, and the skip is counted in the output. + print(f"[dspark-sps] server: pin L={verify_len} refused " + f"({exc}); skipping the rung", file=sys.stderr) + continue + for batch_size in sorted(batch_sizes): + # The load must outlive the measurement window, or the + # cell measures ramp-down rows filed under smaller batch + # sizes. Budget the output by the window at a conservative + # 30 ms/step lower bound. + window_steps = int((settle_s + poll_s) / 0.03) + 64 + max_tokens = max(64, window_steps * (1 + int(verify_len))) + num_requests = int(batch_size) * max(1, int(dp_size)) + prompts = [[rng.randrange(1000, 30000) for _ in range(input_len)] + for _ in range(num_requests)] + bodies = [{"model": model, "prompt": p, "max_tokens": max_tokens, + "temperature": 0.0, "ignore_eos": True, "stream": False} + for p in prompts] + threads = [threading.Thread(target=_fire, args=(b, ), + daemon=True) for b in bodies] + for thread in threads: + thread.start() + time.sleep(settle_s) # let admission reach the plateau + drain_metrics() # discard the ramp + time.sleep(poll_s) + rows = drain_metrics() + # The length is NOT read back from telemetry: /metrics carries + # no generation-token count (iter_states reaches the log line + # only), and it does not need to -- the sweep pinned it, and + # the pin is refused by the server if it could not be honoured. + # The row count is what the step really ran; the padded width + # is resolved from the ladder when one was supplied, since a + # partly-filled batch costs what its captured shape costs. + # + # hostStepTimeMS, not iterLatencyMS: under the overlap + # scheduler the latter spans about two loops (base_worker says + # so where it publishes both), which would inflate every cell + # by roughly a factor of two. + kept = 0 + cell_rows: List[StepSample] = [] + for row in rows: + rows_seen = int(row.get("numActiveRequests") or 0) + step_ms = float(row.get("hostStepTimeMS") or 0.0) + if rows_seen <= 0 or step_ms <= 0: + continue + padded_bs = rows_seen + if padded_batch_sizes: + wider = [int(b) for b in sorted(padded_batch_sizes) + if int(b) >= rows_seen] + if not wider: + continue + padded_bs = wider[0] + cell_rows.append(StepSample(batch_size=padded_bs, + verify_len=int(verify_len), + step_time_ms=step_ms)) + kept += 1 + # Applied, not just printed: the first rows after a pin change + # can straddle the previous window's shape. + out.extend(cell_rows[warmup_steps:]) + note = "" + if not kept: + if post_errors: + note = f" -- requests refused: {post_errors[0]}" + elif not rows: + note = (" -- /metrics returned nothing: is " + "enable_iter_perf_stats on?") + else: + keys = sorted(rows[-1])[:8] + note = (f" -- {len(rows)} stats rows carried no usable " + f"(numActiveRequests, hostStepTimeMS); keys " + f"seen: {keys}") + print(f"[dspark-sps] server bs={batch_size:>5} L={verify_len} " + f"({num_requests} requests in flight) " + f"kept {max(kept - warmup_steps, 0)}" + f" steps{note}", file=sys.stderr) + for thread in threads: + thread.join(timeout=1.0) + finally: + # Never leave a served deployment pinned to a profiling window. + pin(None) + + if server_log: + # Preferred source when the server's log is reachable. /metrics needs + # a metrics collector (return_perf_metrics) to run its tee buffer, and + # its fallback path -- draining the stats queue directly -- came back + # empty on an MPI deployment. The iteration log needs neither: it + # carries the row count, the padded token total and the host step time + # for every step, so the samples are labelled by what ran rather than + # by what was asked for. The pin's job is then only to make sure every + # rung is VISITED. + from_log = samples_from_iter_log( + [server_log], max_draft_len=max_draft_len, + max_batch_size=(max(padded_batch_sizes) if padded_batch_sizes + else 256), + padded_batch_sizes=padded_batch_sizes) + print(f"[dspark-sps] server log yielded {len(from_log)} steps " + f"(/metrics yielded {len(out)})", file=sys.stderr) + if from_log: + return from_log + return out + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = _build_parser().parse_args(argv) + + verify_lens = ( + _int_list(args.verify_lens) + if args.verify_lens + else list(range(1, int(args.max_draft_len) + 1)) + ) + batch_sizes = _int_list(args.batch_sizes) + samples_out = args.samples_out or f"{args.out}.samples.jsonl" + + if args.fit_only or args.samples_in: + if not args.samples_in: + raise SystemExit("--fit-only requires --samples-in") + samples = _load_samples(_str_list(args.samples_in)) + provenance = {"source": "refit", "samples_in": _str_list(args.samples_in)} + # Saved samples may come from any source, including deployment logs + # full of ramp-occupancy cells; require nothing, keep what is thick. + expected_cells = frozenset() + elif args.from_iter_log: + ladder = (_int_list(args.cuda_graph_batch_sizes) + if args.cuda_graph_batch_sizes else None) + samples = samples_from_iter_log( + args.from_iter_log, max_draft_len=args.max_draft_len, + max_batch_size=max(ladder) if ladder else 256, + padded_batch_sizes=ladder) + if not samples: + raise SystemExit( + "no usable steps in the iteration logs. They must come from a " + "run with print_iter_log AND enable_iter_perf_stats enabled: " + "the token total lives in the `states` dict at the end of each " + "line, and without it a step cannot be filed under a cell.") + _write_samples(samples_out, samples) + print(f"[dspark-sps] {len(samples)} steps from " + f"{len(args.from_iter_log)} iteration log(s)", file=sys.stderr) + provenance = {"source": "iter_log", "logs": list(args.from_iter_log)} + # A production log has no requested grid at all. + expected_cells = frozenset() + elif args.base_url: + samples = samples_from_server( + base_url=args.base_url.rstrip("/"), + control_url=(args.control_url.rstrip("/") + if args.control_url else None), + batch_sizes=batch_sizes, + verify_lens=verify_lens, + input_len=args.input_len, + model=args.served_model_name, + warmup_steps=args.warmup_steps, + settle_s=float(args.server_settle_s), + poll_s=float(args.server_poll_s), + max_draft_len=args.max_draft_len, + seed=args.seed, + padded_batch_sizes=(_int_list(args.cuda_graph_batch_sizes) + if args.cuda_graph_batch_sizes else None), + server_log=args.server_log, + # The bs axis is per-rank under attention DP (see the collector's + # docstring); the flags describing the deployment double as the + # fan-out factor. + dp_size=(int(args.tp_size) if args.enable_attention_dp else 1)) + if not samples: + raise SystemExit( + "the server sweep collected nothing. Check that /metrics " + "returns iteration stats (enable_iter_perf_stats) and that " + "/dspark/verify_len_pin exists on this build.") + _write_samples(samples_out, samples) + print(f"[dspark-sps] {len(samples)} steps from {args.base_url}", + file=sys.stderr) + provenance = {"source": "server_sweep", "base_url": args.base_url} + # Which cells the sweep can PROMISE depends on who controls the batch + # size. Without attention DP the server schedules one global batch, so + # a steady client concurrency IS the row count and the requested grid + # is enforceable. Under attention DP the router spreads the load and + # each rank's occupancy is emergent -- 1024 requests over 8 ranks + # oscillates around 128/rank and pads to whatever ladder entry the + # drift lands on (observed: requested 128, measured at 192; only + # the cap-saturated 256 landed exactly). The pin guarantees L; bs is + # whatever the deployment did. Require nothing, keep what is thick, + # and let the fit's grid-connectivity check catch a real hole. + expected_cells = (frozenset( + (int(b), int(l)) for b in batch_sizes for l in verify_lens) + if not args.enable_attention_dp else frozenset()) + else: + if not args.model: + raise SystemExit("--model is required unless --fit-only is used") + extra: Dict[str, object] = {} + if args.extra_llm_api_options: + import yaml + + with open(args.extra_llm_api_options, encoding="utf-8") as handle: + extra = yaml.safe_load(handle) or {} + fracs = _float_list(args.fracs) if args.fracs else None + if fracs: + if args.verify_lens: + raise SystemExit("--fracs and --verify-lens are mutually " + "exclusive; the frac sweep's M is measured, " + "not pinned") + # The frac sweep runs the constant block with live windows; the + # nominal verify_len only labels the samples and sizes the token + # budget, and max_draft_len satisfies every geometry guard. + verify_lens = [int(args.max_draft_len)] + config = SweepConfig( + model=args.model, + speculative_model=args.speculative_model or args.model, + batch_sizes=batch_sizes, + verify_lens=verify_lens, + fracs=fracs, + tp_size=args.tp_size, + ep_size=args.ep_size, + enable_attention_dp=args.enable_attention_dp, + max_draft_len=args.max_draft_len, + input_len=args.input_len, + sts_path=args.sts_path, + seed_table_path=args.seed_table, + stats_poll_s=float(args.stats_poll_s), + block_follows_verify_len=bool(args.block_follows_verify_len), + warmup_steps=args.warmup_steps, + measure_steps=args.measure_steps, + repeats=args.repeats, + min_samples=args.min_samples, + max_seq_len=args.max_seq_len, + max_num_tokens=args.max_num_tokens, + kv_cache_fraction=args.kv_cache_fraction, + attn_backend=args.attn_backend, + disable_overlap_scheduler=args.disable_overlap_scheduler, + enable_confidence_scheduling=args.enable_confidence_scheduling, + timing_key=TIMING_FIELDS[args.timing_field], + pin_acceptance=BONUS_TOKEN_ONLY_ACCEPTANCE if args.pin_acceptance else 0.0, + seed=args.seed, + extra_llm_api_options=extra, + engine_tiers=(_int_list(args.engine_tiers) + if args.engine_tiers else None), + ) + if fracs: + samples = run_frac_sweep(config) + # Which buckets a frac visits is the planner's choice, not a + # promise; thin incidental buckets are dropped by summarize_cells. + expected_cells = None + else: + samples = run_sweep(config) + expected_cells = frozenset( + (int(b), int(l)) for b in batch_sizes for l in verify_lens) + _write_samples(samples_out, samples) + print(f"[dspark-sps] wrote {len(samples)} raw samples to {samples_out}", file=sys.stderr) + provenance = { + "source": "sweep", + "model": config.model, + "tensor_parallel_size": config.tp_size, + "moe_expert_parallel_size": config.ep_size, + "attention_dp": config.enable_attention_dp, + "input_len": config.input_len, + "timing_field": config.timing_key, + "acceptance_pinned": bool(config.pin_acceptance), + "overlap_scheduler": not config.disable_overlap_scheduler, + "confidence_scheduling": config.enable_confidence_scheduling, + # _prepare_environment pins COMPACT for both sweep flavours; the + # pinned sweep's uniform shape comes from the pin, not the mode. + "ragged_verify_mode": RaggedVerifyMode.COMPACT.value, + **({"fracs": [float(f) for f in fracs]} if fracs else {}), + } + + cells = summarize_cells(samples, + # Frac cells discard warmup inside _run_cell (per + # cell, by iteration); trimming again here would + # re-charge it against every (bs, M) bucket. + warmup_steps=(0 if args.fracs else args.warmup_steps), + min_samples=args.min_samples, + expected_cells=expected_cells) + if not cells: + # Distinguished from a short cell: this means no step was ever aligned, + # which is a wiring problem (iteration stats off, the wrong rank count, + # the batch never reaching the requested size) rather than a noisy run. + raise InsufficientSamplesError( + f"no cells at all: {len(samples)} raw samples produced nothing usable. " + f"Check that enable_iter_perf_stats reached the engine, that " + f"--enable-attention-dp matches the deployment (the aligned-step " + f"filter expects one stats row per DP rank), and that the requested " + f"batch sizes were actually reached." + ) + reference_bs = int(args.reference_batch_size or max(c.batch_size for c in cells)) + block_size = max(int(args.max_draft_len), max(c.verify_len for c in cells)) + + payload = build_cost_table_payload( + cells, + riser_tolerance=args.riser_tolerance, + max_breakpoints=args.max_breakpoints, + meta={ + "block_size": block_size, + "min_verify_len": 1, + "modal_batch_size": reference_bs, + "warmup_steps_discarded": int(args.warmup_steps), + "statistic": "median", + # The engine fingerprint the loader verifies (check_table_fingerprint). + # Facts only, and only facts this process can VOUCH for. The loader + # compares every key present in both dicts and refuses on mismatch, + # so a key must be absent -- not defaulted -- when the fact is + # unknown: recording ep as `args.ep_size or 0` fabricated ep=0 + # whenever the flag was unset, the loader's Mapping normalizes + # unset to 1, and the profiler's own table was refused on the very + # engine that produced it. The deployment-measuring sources + # (--base-url, --from-iter-log) build no engine here at all, so + # they attach no engine dict; the deployment's own config is the + # fingerprint authority there, and an unfingerprinted table loads + # with a warning rather than a false claim. + **({"engine": { + **({"tp": int(args.tp_size)} + if args.tp_size is not None else {}), + **({"ep": int(args.ep_size)} + if args.ep_size else {}), + "attention_dp": bool(args.enable_attention_dp), + "block": block_size, + "max_batch_size": max(_int_list(args.batch_sizes)), + "geometry": ("block_follows_verify_len" + if getattr(args, "block_follows_verify_len", False) + else "constant_block"), + }} if provenance.get("source") == "sweep" else {}), + **provenance, + }, + ) + + tiers = ( + _int_list(args.verify_len_tiers) + if args.verify_len_tiers + else derive_verify_len_tiers( + cost_table=load_cost_table(payload), + num_requests=reference_bs, + block_size=block_size, + min_verify_len=1, + max_tiers=int(args.max_tiers), + ) + ) + payload["_meta"]["recommended_confidence_verify_len_tiers"] = [int(t) for t in tiers] + payload["_meta"]["recommended_tiers_note"] = ( + "Nothing reads this from the table file. Copy it into " + "DSparkDecodingConfig.confidence_verify_len_tiers, which is also what " + "sizes the captured CUDA graph set." + ) + + diagnostics = check_table_is_informative( + payload, + batch_sizes=sorted({c.batch_size for c in cells}), + tiers=tiers, + acceptance_rates=_float_list(args.probe_acceptance_rates), + min_step_time_spread=args.min_step_time_spread, + allow_inert=args.allow_inert, + ) + payload["_meta"]["diagnostics"] = diagnostics + + with open(args.out, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=False) + handle.write("\n") + + _print_report(payload, cells, diagnostics, tiers) + return 0 + + +def _str_list(text: str) -> List[str]: + return [v for v in str(text).replace(",", " ").split() if v] + + +def _print_report( + payload: dict, + cells: Sequence[CellStat], + diagnostics: dict, + tiers: Sequence[int], +) -> None: + meta = payload["_meta"] + print("") + print("DSpark SPS cost table") + print("---------------------") + print(f" token_counts : {payload['token_counts']}") + print(f" step_time_ms : {payload['step_time_ms']} (theta(M) only)") + print(f" fixed_overhead_ms : {payload['fixed_overhead_ms']}") + print(f" batch_sizes : {payload['batch_sizes']}") + print(f" batch_overhead_ms : {payload['batch_overhead_ms']}") + print("") + print(f" cells measured : {len(cells)}") + print( + f" additive-model fit: max residual " + f"{meta['fit']['max_abs_residual_ms']:.4f} ms " + f"({meta['fit']['max_rel_residual']:.2%} of a typical step)" + ) + for warning in meta["fit"]["warnings"]: + print(f" WARNING: {warning}") + print(f" step-time spread : {diagnostics['best_step_time_spread']:.2%} across the ladder") + print("") + print(f" Copy into confidence_verify_len_tiers: {list(tiers)}") + print(" Planner choice vs synthetic acceptance rate (batch_size / rate -> length):") + for probe in diagnostics["profitability_probe"]: + mark = "trim" if probe["trims"] else " " + print( + f" bs={probe['batch_size']:>5} p={probe['acceptance_rate']:.2f} " + f"-> L={probe['chosen_verify_len']} (max {probe['max_tier']}) {mark}" + ) + print("") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/microbenchmarks/dsv4_sparse_mla_presentation.py b/tests/microbenchmarks/dsv4_sparse_mla_presentation.py new file mode 100644 index 000000000000..4ab64673ad92 --- /dev/null +++ b/tests/microbenchmarks/dsv4_sparse_mla_presentation.py @@ -0,0 +1,457 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""B2: batch-major vs token-major presentation of DeepSeek-V4 sparse-MLA generation. + +This is the measurement the DSpark ragged-verification design hinges on. + +The question +------------ +Ragged verification means each request sends its own number of query tokens to +the target in one decode step. Every generation kernel that recovers ``(request, +offset)`` from a single ``next_n`` breaks under that. The deepest one is +trtllm-gen sparse-MLA generation, which cannot be read: its kernel selection +lives in a precompiled cubin, and the host side asserts a uniform layout:: + + TLLM_CHECK_WITH_INFO(num_tokens % num_seqs == 0, + "seq_len should be same for all generation requests", ...) + int32_t const input_seq_length = num_tokens / num_seqs; + -- cpp/tensorrt_llm/thop/attentionOp.cpp + mMaxSeqLenQ = acc_q_len / batch_beam; + -- cpp/tensorrt_llm/common/attentionOp.cpp + +Both candidate ragged designs (D1 token-flat, D2 token-major) dodge that assert +the same way: present the generation half as ``num_gen_tokens`` sequences of +length one, so ``seq_len == 1`` and the check passes trivially. They are +therefore the same bet, and this script is how that bet gets priced. + +What is compared +---------------- +Identical total query tokens and an identical per-token sparse index budget, +laid out two ways: + +* **batch-major** -- ``B`` requests, ``s_q = tier + 1``, each cached at ``L``. + This is what the engine does today. +* **token-major** -- ``B * (tier + 1)`` requests, ``s_q = 1``, request ``(i, j)`` + cached at ``L + j``. Each query token is its own sequence, so it carries its + own KV length and block table and no kernel has to divide anything. + +Row ``j`` of request ``i`` attends over ``L + j + 1`` tokens in both layouts, so +the arithmetic is the same; only the presentation differs. + +How to read the result +---------------------- +If token-major is materially slower or incorrect, the fallback is the uniform +tier ladder, which needs no kernel change at all. The prior is that DeepSeek-V4-Pro under DEP8 +keeps all 128 Q heads per rank, so BMM1's M dimension falls from +``128 * (tier+1)`` to ``128`` -- and because the sparse index set differs per +row there was never any KV tile sharing across the window to lose. That is an +inference, which is exactly why it needs measuring. + +Scope: this times the whole attention layer, so it includes the indexer +(paged-MQA-logits + top-k), which itself changes presentation -- the batch-major +layout takes the expanded one-row-per-token path, the token-major layout takes +the strided ``next_n == 1`` path. That bundling is deliberate: it is what a +ragged step would actually pay, and it subsumes B3. The pure-attention delta +cannot be separated without a kernel-level profiler. + +Usage (single GPU, no model weights):: + + python tests/microbenchmarks/dsv4_sparse_mla_presentation.py + python tests/microbenchmarks/dsv4_sparse_mla_presentation.py \ + --batch-sizes 8 32 128 --tiers 1 3 5 --context-len 4096 --check +""" + +import argparse +import json +import math +import os +import sys +from typing import Dict, List, Optional, Tuple + +import torch + +TESTS_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +# Reuse the sparse-MLA test harness verbatim rather than restating it: it is the +# only place that builds a real DeepseekV4TrtllmAttention plus cache manager, +# and a second copy would drift. +sys.path.insert(0, os.path.join(TESTS_ROOT, "unittest", "_torch", "attention", + "sparse", "deepseek_v4")) + +from test_deepseek_v4_sparse_mla import ( # noqa: E402 + Scenario, _build_compressed_topk_indices, _create_cache_manager, + _create_pos_embd_params, _create_rope_cos_sin, _prefill_compress_buffer) + +from tensorrt_llm._torch.attention_backend.interface import ( # noqa: E402 + AttentionInputType, MLAParams) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import \ + DeepseekV4TrtllmAttention # noqa: E402 +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import \ + DeepseekV4TrtllmAttentionMetadata # noqa: E402 +from tensorrt_llm._torch.metadata import KVCacheParams # noqa: E402 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest # noqa: E402 +from tensorrt_llm.bindings import SamplingConfig # noqa: E402 +from tensorrt_llm.mapping import Mapping # noqa: E402 + +# Layer 1 of the shared Scenario is the compress_ratio=4 layer, i.e. the one +# that actually runs sparse selection. Layers 0 (dense) and 2 (ratio 128) do not +# exercise the path this question is about. +SPARSE_LAYER_IDX = 1 + + +def _yarn_mscale(scale: float = 1.0, mscale: float = 1.0) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class PresentationHarness: + """One cache manager + attention layer, driven at a chosen presentation.""" + + def __init__(self, scenario: Scenario, cached_lens: List[int], + seq_len_q: int, device: torch.device): + """``cached_lens[i]`` is request i's KV length before this step; every + request contributes ``seq_len_q`` query tokens.""" + self.scenario = scenario + self.device = device + self.cached_lens = list(cached_lens) + self.seq_len_q = seq_len_q + self.num_requests = len(cached_lens) + self.total_q_tokens = self.num_requests * seq_len_q + + self.qk_rope_head_dim = scenario.qk_rope_head_dim + # rope_append=False: the last 64 of the 512 lanes carry the rope part, + # and the head count halves. Mirrors the test harness. + self.kv_lora_rank = scenario.kv_lora_rank - self.qk_rope_head_dim + self.head_dim = self.kv_lora_rank + self.qk_rope_head_dim + self.num_heads = 64 + + # The pool is sized `max_seq_len * num_requests`, but allocation is by + # whole blocks, so each request really consumes + # `ceil(len / block) * block`. Round up and add a block of slack, or the + # token-major layout -- which has `next_n` times as many requests, and so + # pays the rounding `next_n` times as often -- runs the pool dry and the + # cache manager refuses the context before it is ever measured. + block = scenario.kv_cache_tokens_per_block + raw_max_seq_len = max(cached_lens) + seq_len_q + 1 + max_seq_len = ((raw_max_seq_len + block - 1) // block + 1) * block + # The cache manager sizes its pool from the *context* lengths it is + # handed, so pass the per-request KV lengths as if they were prompts. + self.cache_manager, self.sparse_config = _create_cache_manager( + scenario, self.cached_lens, max_seq_len) + self.request_ids = list(range(self.num_requests)) + self.requests = [ + LlmRequest( + request_id=i, + max_new_tokens=seq_len_q + 1, + input_tokens=list(range(cached_len)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) for i, cached_len in enumerate(self.cached_lens) + ] + for request in self.requests: + self.cache_manager.prepare_context(request) + self.cache_manager.resize_context(request, + request.context_chunk_size) + # The context allocation covers the cached tokens only. Each query token + # this step appends one more, which can cross a page boundary and needs + # a block the block table does not have yet -- reading it is an illegal + # access, not a graceful failure. + for _ in range(seq_len_q): + for request in self.requests: + assert self.cache_manager.try_allocate_generation(request), ( + f"could not extend request {request.py_request_id} for " + f"generation") + + self.mapping = Mapping(world_size=1, tp_size=1, rank=0) + self.rope_cos_sin = _create_rope_cos_sin(scenario, device) + pos_embd_params = _create_pos_embd_params(scenario) + mscale = _yarn_mscale(pos_embd_params.rope.scale, + pos_embd_params.rope.mscale_all_dim) + q_scaling = 1.0 / (mscale * mscale) + + self.layer = DeepseekV4TrtllmAttention( + layer_idx=SPARSE_LAYER_IDX, + num_heads=self.num_heads, + head_dim=self.head_dim, + num_kv_heads=1, + q_scaling=q_scaling, + pos_embd_params=pos_embd_params, + mla_params=MLAParams( + q_lora_rank=scenario.q_lora_rank, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + qk_nope_head_dim=scenario.qk_nope_head_dim, + v_head_dim=scenario.v_head_dim, + rope_append=scenario.rope_append, + predicted_tokens_per_seq=1, + hidden_size=scenario.hidden_size, + ), + sparse_attention_config=self.sparse_config, + skip_create_weights_in_init=True, + ) + self.layer.update_quant_config(None) + self.layer.attn_sink = torch.nn.Parameter(torch.randn( + self.num_heads, dtype=torch.float32, device=device).mul_(0.5), + requires_grad=False) + + _prefill_compress_buffer(self.cache_manager, SPARSE_LAYER_IDX, + self.cached_lens, self.request_ids, + self.head_dim, device) + + self._build_inputs() + self._build_metadata() + + def _build_inputs(self) -> None: + dtype = self.scenario.dtype + device = self.device + n = self.total_q_tokens + self.compressed_kv = torch.empty([n, self.kv_lora_rank], + dtype=dtype, + device=device).uniform_(-1, 1) + self.k_pe = torch.empty([n, self.qk_rope_head_dim], + dtype=dtype, + device=device).uniform_(-1, 1) + q = torch.empty([n, self.num_heads, self.kv_lora_rank], + dtype=dtype, + device=device).uniform_(-1, 1) + self.q_pe = torch.empty([n, self.num_heads, self.qk_rope_head_dim], + dtype=dtype, + device=device).uniform_(-1, 1) + self.fused_q = torch.cat([q, self.q_pe], + dim=-1).view(-1, + self.num_heads * self.head_dim) + self.latent_cache = torch.cat([self.compressed_kv, self.k_pe], dim=-1) + + def _build_metadata(self) -> None: + device = self.device + self.metadata = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor([self.seq_len_q] * self.num_requests, + dtype=torch.int), + request_ids=self.request_ids, + max_num_requests=self.num_requests, + num_contexts=0, + prompt_lens=self.cached_lens, + max_num_tokens=self.total_q_tokens, + kv_cache_manager=self.cache_manager, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=self.cached_lens, + ), + mapping=self.mapping, + enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), + sparse_attention_config=self.sparse_config, + ) + self.metadata.prepare() + + num_seqs = self.metadata.kv_lens_cuda_runtime.size(0) + self.cu_q_seqlens = torch.empty(num_seqs + 1, + dtype=torch.int32, + device=device) + self.cu_kv_seqlens = torch.empty(num_seqs + 1, + dtype=torch.int32, + device=device) + self.fmha_scheduler_counter = torch.empty(1, + dtype=torch.uint32, + device=device) + # Per-query-token KV extent, i.e. exactly what a correct ragged layout + # must produce: row j of a request cached at L sees L + j + 1 tokens. + per_row_kv_lens = [ + cached_len + j + 1 for cached_len in self.cached_lens + for j in range(self.seq_len_q) + ] + self.topk_indices = _build_compressed_topk_indices( + [kv_len - 1 for kv_len in per_row_kv_lens], + self.scenario.compress_ratios[SPARSE_LAYER_IDX], + self.scenario.index_topk, + device, + ) + + def step(self) -> torch.Tensor: + """One generation step. Repeatable: every write target comes from the + metadata, which is fixed, so re-running rewrites the same slots.""" + self.layer.mla_rope_generation( + self.fused_q, + self.q_pe, + self.latent_cache, + self.metadata, + self.cu_q_seqlens, + self.cu_kv_seqlens, + self.fmha_scheduler_counter, + None, + None, + None, + ) + return self.layer.forward( + self.fused_q, + None, + None, + self.metadata, + attention_input_type=AttentionInputType.generation_only, + latent_cache=self.latent_cache, + q_pe=self.q_pe, + cu_q_seqlens=self.cu_q_seqlens, + cu_kv_seqlens=self.cu_kv_seqlens, + fmha_scheduler_counter=self.fmha_scheduler_counter, + topk_indices=self.topk_indices, + ) + + def shutdown(self) -> None: + self.cache_manager.shutdown() + + +def _time_step(harness: PresentationHarness, warmup: int, + iters: int) -> Tuple[float, float]: + """Median and minimum per-step latency in milliseconds. + + Median rather than mean: a single stray context switch on a shared node + should not decide a design. + """ + for _ in range(warmup): + harness.step() + torch.cuda.synchronize() + + samples: List[float] = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + harness.step() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end)) + samples.sort() + return samples[len(samples) // 2], samples[0] + + +def _run_point(batch_size: int, tier: int, context_len: int, warmup: int, + iters: int, check: bool) -> Dict[str, object]: + next_n = tier + 1 + scenario = Scenario() + device = torch.device("cuda") + torch.manual_seed(42) + + record: Dict[str, object] = { + "batch_size": batch_size, + "tier": tier, + "next_n": next_n, + "context_len": context_len, + "total_query_tokens": batch_size * next_n, + } + + # Batch-major: B sequences, next_n query tokens each. + batch_major = PresentationHarness(scenario, [context_len] * batch_size, + next_n, device) + # Token-major: one sequence per query token. Request (i, j) is cached at + # L + j so its single query token has the same extent as row j above. + token_cached = [ + context_len + j for _ in range(batch_size) for j in range(next_n) + ] + token_major = PresentationHarness(scenario, token_cached, 1, device) + + try: + if check: + out_batch = batch_major.step() + out_token = token_major.step() + record["batch_major_output_shape"] = list(out_batch.shape) + record["token_major_output_shape"] = list(out_token.shape) + record["shapes_match"] = (out_batch.shape == out_token.shape) + for name, out in (("batch_major", out_batch), ("token_major", + out_token)): + record[f"{name}_finite"] = bool(torch.isfinite(out).all()) + record[f"{name}_abs_mean"] = float(out.abs().mean()) + # The two runs use different random KV, so the outputs are not + # expected to match elementwise -- only to be the same shape, finite, + # and of the same magnitude. A token-major layout that mis-attributes + # rows shows up here as a magnitude that is off, or as non-finite + # values, because each pseudo-request's extent differs by one. + + batch_median, batch_min = _time_step(batch_major, warmup, iters) + token_median, token_min = _time_step(token_major, warmup, iters) + finally: + batch_major.shutdown() + token_major.shutdown() + + record["batch_major_ms_median"] = batch_median + record["batch_major_ms_min"] = batch_min + record["token_major_ms_median"] = token_median + record["token_major_ms_min"] = token_min + record["token_major_slowdown"] = (token_median / + batch_median if batch_median else + float("inf")) + return record + + +def main(argv: Optional[List[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--batch-sizes", + type=int, + nargs="+", + default=[8, 32, 128]) + parser.add_argument("--tiers", + type=int, + nargs="+", + default=[1, 3, 5], + help="verify-length tiers; next_n = tier + 1") + parser.add_argument("--context-len", type=int, default=4096) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=30) + parser.add_argument("--check", + action="store_true", + help="also run a one-shot sanity check per point") + parser.add_argument("--json", type=str, default=None) + args = parser.parse_args(argv) + + if not torch.cuda.is_available(): + print("CUDA is required for this benchmark", file=sys.stderr) + return 2 + + print(f"device: {torch.cuda.get_device_name(0)} " + f"(sm{''.join(str(c) for c in torch.cuda.get_device_capability())})") + print(f"context_len={args.context_len} warmup={args.warmup} " + f"iters={args.iters}\n") + header = (f"{'B':>5} {'tier':>5} {'next_n':>7} {'tokens':>7} " + f"{'batch-major ms':>15} {'token-major ms':>15} {'slowdown':>9}") + print(header) + print("-" * len(header)) + + records: List[Dict[str, object]] = [] + for batch_size in args.batch_sizes: + for tier in args.tiers: + try: + record = _run_point(batch_size, tier, args.context_len, + args.warmup, args.iters, args.check) + except Exception as exc: # noqa: BLE001 - report and keep going + print(f"{batch_size:>5} {tier:>5} {'':>7} {'':>7} " + f"FAILED: {type(exc).__name__}: {exc}") + records.append({ + "batch_size": batch_size, + "tier": tier, + "error": f"{type(exc).__name__}: {exc}", + }) + continue + records.append(record) + print(f"{record['batch_size']:>5} {record['tier']:>5} " + f"{record['next_n']:>7} {record['total_query_tokens']:>7} " + f"{record['batch_major_ms_median']:>15.3f} " + f"{record['token_major_ms_median']:>15.3f} " + f"{record['token_major_slowdown']:>9.2f}x") + + ok = [r for r in records if "error" not in r] + if ok: + worst = max(r["token_major_slowdown"] for r in ok) + best = min(r["token_major_slowdown"] for r in ok) + print(f"\ntoken-major slowdown across points: " + f"{best:.2f}x .. {worst:.2f}x") + print("Read: >~1.3x sustained means the ragged (D1/D2) route is paying " + "more than it can save; fall back to the uniform tier ladder " + "(D3 Stage A), which needs no kernel change.") + + if args.json: + with open(args.json, "w") as handle: + json.dump(records, handle, indent=2) + print(f"wrote {args.json}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py index e7df3560db53..4c40e44e51cf 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py @@ -18,7 +18,6 @@ """ import math -import os from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple @@ -27,7 +26,6 @@ from utils.util import skip_pre_blackwell from tensorrt_llm._torch.attention_backend.interface import ( - AttentionForwardArgs, AttentionInputType, MLAParams, PositionalEmbeddingParams, @@ -40,7 +38,6 @@ ) from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( DeepseekV4TrtllmAttentionMetadata, - get_token_bytes, ) from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest @@ -106,6 +103,11 @@ class RopeConfig: # Layers to test: layer 0 → compress_ratio=1, layer 1 → ratio=4, layer 2 → ratio=128 TEST_LAYERS = [0, 1, 2] +# The widest generation window the parameterized test exercises, i.e. the static +# ceiling the DSA metadata sizes its expanded buffers from. Matches +# DeepSeek-V4-Pro-DSpark's max_draft_len=5 -> next_n=6. +MAX_GENERATION_SEQ_LEN_Q = 6 + # RoPE helpers def rotate_half(x): @@ -399,6 +401,33 @@ def calculate_deepseek_v4_ref_ctx_sparse( return torch.cat(ref_results, dim=0) +def _seq_lens_q_or_uniform( + num_query_tokens: int, num_requests: int, seq_lens_q: Optional[List[int]] +) -> List[int]: + """Per-request query-token counts, defaulting to the uniform split. + + The reference implementations below used to derive this as + ``num_query_tokens // num_requests``. That is the assumption under test once + a scheduler can hand each request its own verify window, so it is a + parameter now; passing None keeps the historical behavior for callers whose + batches really are uniform. + """ + if seq_lens_q is not None: + assert len(seq_lens_q) == num_requests, ( + f"got {len(seq_lens_q)} query lengths for {num_requests} requests" + ) + assert sum(seq_lens_q) == num_query_tokens, ( + f"query lengths sum to {sum(seq_lens_q)} but {num_query_tokens} " + f"query tokens were supplied" + ) + return [int(v) for v in seq_lens_q] + assert num_query_tokens % num_requests == 0, ( + f"{num_query_tokens} query tokens do not split evenly across " + f"{num_requests} requests; pass seq_lens_q explicitly" + ) + return [num_query_tokens // num_requests] * num_requests + + def _rotate_gen_inputs( fused_q: torch.Tensor, q_pe: torch.Tensor, @@ -409,23 +438,29 @@ def _rotate_gen_inputs( num_heads: int, kv_lora_rank: int, qk_rope_head_dim: int, + seq_lens_q: Optional[List[int]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: fused_head_dim = kv_lora_rank + qk_rope_head_dim num_requests = len(seq_lens_kv) - seq_len_q = fused_q.shape[0] // num_requests + q_lens = _seq_lens_q_or_uniform(fused_q.shape[0], num_requests, seq_lens_q) + q_starts = [0] + for q_len in q_lens: + q_starts.append(q_starts[-1] + q_len) fused_q_rot = fused_q.clone() new_latent_list = [] for i in range(num_requests): past_len = seq_lens_kv[i] + seq_len_q = q_lens[i] + lo, hi = q_starts[i], q_starts[i + 1] - fused_q_seq = fused_q_rot[i * seq_len_q : (i + 1) * seq_len_q].unflatten( + fused_q_seq = fused_q_rot[lo:hi].unflatten( -1, [num_heads, fused_head_dim] ) - q_pe_seq = q_pe[i * seq_len_q : (i + 1) * seq_len_q] - compressed_kv_seq = compressed_kv[i * seq_len_q : (i + 1) * seq_len_q] - k_pe_seq = k_pe[i * seq_len_q : (i + 1) * seq_len_q].unsqueeze(-2) + q_pe_seq = q_pe[lo:hi] + compressed_kv_seq = compressed_kv[lo:hi] + k_pe_seq = k_pe[lo:hi].unsqueeze(-2) cos, sin = rope_cos_sin[past_len : past_len + seq_len_q].chunk(2, dim=-2) @@ -465,25 +500,33 @@ def calculate_deepseek_v4_ref_gen_sparse( q_scaling: float, compress_ratio: int, attn_sink: Optional[torch.Tensor] = None, + seq_lens_q: Optional[List[int]] = None, ): """Reference attention for DeepSeek-V4 generation phase.""" fused_head_dim = kv_lora_rank + qk_rope_head_dim bmm1_scale = 1 / (math.sqrt(qk_nope_head_dim + qk_rope_head_dim) * q_scaling) num_requests = len(seq_lens_kv) - seq_len_q = fused_q_rot.shape[0] // num_requests + q_lens = _seq_lens_q_or_uniform( + fused_q_rot.shape[0], num_requests, seq_lens_q + ) + q_starts = [0] + for q_len in q_lens: + q_starts.append(q_starts[-1] + q_len) ref_results = [] latent_cache_list = [] total_past_tokens = 0 for i in range(num_requests): past_len = seq_lens_kv[i] + seq_len_q = q_lens[i] + lo, hi = q_starts[i], q_starts[i + 1] - fused_q_seq = fused_q_rot[i * seq_len_q : (i + 1) * seq_len_q].unflatten( + fused_q_seq = fused_q_rot[lo:hi].unflatten( -1, [num_heads, fused_head_dim] ) # New token's latent cache - new_token_latent = new_latent_cache[i * seq_len_q : (i + 1) * seq_len_q].unsqueeze( + new_token_latent = new_latent_cache[lo:hi].unsqueeze( -2 ) # [seq_len_q, 1, head_dim] @@ -510,7 +553,7 @@ def calculate_deepseek_v4_ref_gen_sparse( # Compressed KV if compress_ratio > 1 and compressed_ref_data is not None: if compress_ratio == 4 and compressed_topk_indices is not None: - row = i * seq_len_q + qi + row = lo + qi indices_row = compressed_topk_indices[row] valid = indices_row[indices_row >= 0] comp_kv = compressed_ref_data[i][valid.long()] @@ -547,11 +590,21 @@ def calculate_deepseek_v4_ref_gen_sparse( return ref_result, new_latent_cache_out -def _allocate_kv_cache_for_generation(cache_manager, requests: List[LlmRequest]): - for req in requests: - assert cache_manager.try_allocate_generation(req), ( - f"Failed to allocate generation KV cache for request {req.py_request_id}" - ) +def _allocate_kv_cache_for_generation( + cache_manager, requests: List[LlmRequest], num_new_tokens: int = 1 +): + """Extend each request's KV allocation by ``num_new_tokens``. + + One call covers one token. A speculative step appends `1 + draft_len` of + them, and the shortfall only bites once a request crosses a page boundary -- + at which point the block table has no entry for the new block and the read + is an illegal access rather than a clean failure. + """ + for _ in range(num_new_tokens): + for req in requests: + assert cache_manager.try_allocate_generation(req), ( + f"Failed to allocate generation KV cache for request {req.py_request_id}" + ) def _create_rope_config(scenario: Scenario) -> RopeConfig: @@ -620,7 +673,7 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): torch.manual_seed(42) context_lengths = [1] - local_num_heads = int(os.environ.get("DSV4_REPRO_LOCAL_NUM_HEADS", "16")) + local_num_heads = 16 layer_idx = 0 ratio = scenario.compress_ratios[layer_idx] assert ratio == 1 @@ -671,8 +724,7 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): ) layer.update_quant_config(None) attn_sink = torch.randn(local_num_heads, dtype=torch.float32, device=device).mul_(0.5) - if not os.environ.get("DSV4_REPRO_NO_SINK"): - layer.attn_sink = torch.nn.Parameter(attn_sink, requires_grad=False) + layer.attn_sink = torch.nn.Parameter(attn_sink, requires_grad=False) attn_metadata = DeepseekV4TrtllmAttentionMetadata( seq_lens=torch.tensor(context_lengths, dtype=torch.int), @@ -703,75 +755,7 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): fused_q = torch.cat([q_nope, q_pe], dim=-1).view(1, local_num_heads * head_dim) latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) - simple_swa_pool = None - if os.environ.get("DSV4_REPRO_SIMPLE_POOL"): - simple_swa_pool = torch.empty( - (cache_manager.tokens_per_block, head_dim), dtype=dtype, device=device - ) - simple_swa_pool.zero_() - simple_swa_pool_ptr = simple_swa_pool.data_ptr() - cache_manager.kv_cache_pool_pointers[0, 0] = simple_swa_pool_ptr - attn_metadata.host_kv_cache_pool_pointers[0, 0] = simple_swa_pool_ptr - attn_metadata.sparse_mla_base_ptrs[1] = simple_swa_pool_ptr - attn_metadata.swa_buffer_ptrs[layer_idx] = simple_swa_pool_ptr - attn_metadata.block_tables[(1, DeepseekV4AttentionType.SWA)][0].fill_(-1) - attn_metadata.block_tables[(1, DeepseekV4AttentionType.SWA)][0, 0] = 0 - torch.cuda.synchronize() - - if os.environ.get("DSV4_REPRO_PRINT_PARAMS"): - sparse_attn_indices, _ = layer.sparse_attn_predict( - fused_q, - None, - attn_metadata, - AttentionForwardArgs( - attention_input_type=AttentionInputType.context_only, - latent_cache=latent_cache, - q_pe=q_pe, - ), - ) - torch.cuda.synchronize() - print("DSV4_REPRO local_num_heads", local_num_heads) - print( - "DSV4_REPRO token_stride", - get_token_bytes( - head_dim, - sparse_config.index_head_dim, - ratio, - DeepseekV4AttentionType.SWA, - False, - ), - ) - print("DSV4_REPRO tokens_per_block", cache_manager.tokens_per_block) - print("DSV4_REPRO max_blocks_per_seq", cache_manager.max_blocks_per_seq) - print("DSV4_REPRO max_seq_len", cache_manager.max_seq_len) - print("DSV4_REPRO metadata_max_seq_len", attn_metadata.max_seq_len) - print("DSV4_REPRO kv_lens_runtime", attn_metadata.kv_lens_runtime[:1].cpu().tolist()) - print("DSV4_REPRO host_total_kv_lens", attn_metadata.host_total_kv_lens.cpu().tolist()) - print( - "DSV4_REPRO prompt_lens_runtime", - attn_metadata.prompt_lens_cpu_runtime[:1].cpu().tolist(), - ) - print("DSV4_REPRO swa_pool_base_ptr", attn_metadata.sparse_mla_base_ptrs[1]) - print("DSV4_REPRO swa_buffer_ptr", attn_metadata.swa_buffer_ptrs[layer_idx]) - print( - "DSV4_REPRO block_table_swa", - attn_metadata.block_tables[(1, DeepseekV4AttentionType.SWA)][:1, :4].cpu().tolist(), - ) - print("DSV4_REPRO sparse_attn_indices", sparse_attn_indices.cpu().tolist()) - print("DSV4_REPRO sparse_attn_indices_dtype", sparse_attn_indices.dtype) - print( - "DSV4_REPRO sparse_mla_topk_lens", - attn_metadata.sparse_mla_topk_lens[ratio][:1].cpu().tolist(), - ) - print( - "DSV4_REPRO sparse_mla_topk_lens_dtype", attn_metadata.sparse_mla_topk_lens[ratio].dtype - ) - softmax_stats_tensor = None - if os.environ.get("DSV4_REPRO_SOFTMAX_STATS"): - softmax_stats_tensor = torch.empty( - (1, local_num_heads, 2), dtype=torch.float32, device=device - ) result = layer.forward( fused_q.clone(), @@ -821,8 +805,27 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): @pytest.mark.skip_less_device_memory(80000) @pytest.mark.parametrize("context_lengths", [[4399], [14, 508, 3947], [2, 1406, 3327]]) @pytest.mark.parametrize("num_generation_steps", [2]) -def test_deepseek_v4_sparse_mla(context_lengths: List[int], num_generation_steps: int): - generation_seq_len_q = 1 +# The query width per generation request. Speculative decoding makes this the +# batch-wide `1 + draft_len`, and DSpark's confidence scheduler moves it between +# steps along a captured tier ladder -- `[1, 3, 5]` for max_draft_len=5, i.e. +# widths 2, 4 and 6. Pinning it at 1 left every width-dependent stride in the +# DSA generation path untested: buffers are laid out at the static ceiling while +# the kernels are launched at the runtime width, and when the two disagree the +# result is either misattributed rows or a hang, never an error. +# +# Coverage limit, stated because it is not obvious and it matters: this test +# supplies `topk_indices` directly for the compress_ratio=4 layer, so +# `Indexer.forward` -- the paged-MQA-logits plus top-k path that consumes +# `kv_lens_expanded` / `block_table_expanded` / the DeepGEMM schedule -- never +# runs here. Those buffers are exactly where a wrong stride misattributes rows, +# so widening this parameterization does NOT by itself cover that. Verified +# empirically: reverting the stride fix leaves all twelve cases green. What this +# does cover is the RoPE, KV-write, sparse-MLA and compressor paths at widths +# other than one, which was previously nothing at all. +@pytest.mark.parametrize("generation_seq_len_q", [1, 2, 6]) +def test_deepseek_v4_sparse_mla( + context_lengths: List[int], num_generation_steps: int, generation_seq_len_q: int +): scenario = Scenario() device = torch.device("cuda") dtype = scenario.dtype @@ -1146,7 +1149,9 @@ def yarn_get_mscale(scale=1, mscale=1): # 8. Generation steps for step in range(num_generation_steps): print(f"\n=== Generation step {step + 1} ===") - _allocate_kv_cache_for_generation(cache_manager, requests) + _allocate_kv_cache_for_generation( + cache_manager, requests, num_new_tokens=generation_seq_len_q + ) cached_lens = [ctx_len + step * generation_seq_len_q for ctx_len in context_lengths] kv_lens = [cl + generation_seq_len_q for cl in cached_lens] @@ -1184,6 +1189,21 @@ def yarn_get_mscale(scale=1, mscale=1): enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), sparse_attention_config=sparse_config, ) + # Reproduce the split the engine actually runs with rather than sizing + # everything to the current width. `max_draft_tokens` is the static + # ceiling: the engine hands it `tokens_per_gen_step - 1` once (via + # update_spec_dec_param) and never moves it, and every expanded buffer + # is sized from it so a shorter step cannot trigger a reallocation + # mid-capture. The per-step width arrives separately, through + # `runtime_tokens_per_gen_step`. Letting the two differ is the whole + # point: when they were conflated, a step narrower than the ceiling + # wrote its kv_lens and block table one stride apart and read them back + # at another. Set post-construction and rebuild the width-dependent + # buffers, which is exactly what update_spec_dec_param does. + gen_metadata.max_draft_tokens = MAX_GENERATION_SEQ_LEN_Q - 1 + gen_metadata.runtime_tokens_per_gen_step = generation_seq_len_q + gen_metadata._create_kv_lens_2d_buffer() + gen_metadata.create_expanded_buffers() gen_metadata.prepare() for layer_idx in TEST_LAYERS: @@ -1216,11 +1236,19 @@ def yarn_get_mscale(scale=1, mscale=1): None, # quant_q_buffer ) - # Build topk_indices for ratio=4 + # Build topk_indices for ratio=4. One row per *query token*, not + # per request: query token j of a request cached at L sits at + # position L + j and so sees a different number of compressed + # tokens. Collapsing that to one row per request is the same + # uniform-window assumption this test is parameterized to break. topk_indices = None if ratio == 4: topk_indices = _build_compressed_topk_indices( - [kv - 1 for kv in kv_lens], + [ + cached_len + j + for cached_len in cached_lens + for j in range(generation_seq_len_q) + ], ratio, scenario.index_topk, device, diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py new file mode 100644 index 000000000000..231673b8e0a9 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Token-major row expansion for the DSv4 MLA generation ops. + +The kernel-level tests in ``test_deepseek_v4_sparse_mla.py`` call the ops with +hand-built tensors, so they pass whatever the metadata layer produces -- they +cannot see a view that is assembled wrongly. That gap let a block table shaped +``[rows, max_blocks]`` reach an op that reads +``[num_pools, num_seqs, 2, max_blocks]``: the gather collapsed the leading axes +until only two were left, which stripped ``num_seqs`` rather than the pool axis +and left request ids indexing a length-2 K/V axis. It surfaced only on an +8-GPU end-to-end run, as a device-side ``vectorized gather kernel index out of +bounds``. + +These tests exercise the expansion itself against a stub cache manager, so the +same mistake fails in seconds on one GPU. +""" + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.sparse.dsa import ( + DSAtrtllmAttentionMetadata) + + +class _StubCacheManager: + """Only the attributes the row expansion reads.""" + + def __init__(self, num_pools: int, max_blocks_per_seq: int): + self.num_pools = num_pools + self.num_attention_op_pools = num_pools + self.max_blocks_per_seq = max_blocks_per_seq + + +def _metadata(*, num_pools, max_blocks, max_seqs, draft=5): + """A metadata object with just enough state to build its own buffers. + + Constructed through ``__new__`` rather than the real constructor: the real + one wants a runtime, a resource manager and a model config, none of which + the row expansion touches. What it does touch is set explicitly below, so a + new dependency shows up as a clear AttributeError here instead of being + silently satisfied by a mock. + """ + meta = DSAtrtllmAttentionMetadata.__new__(DSAtrtllmAttentionMetadata) + meta.kv_cache_manager = _StubCacheManager(num_pools, max_blocks) + meta.max_num_sequences = max_seqs + meta.max_num_requests = max_seqs + meta.max_draft_tokens = draft + meta.cuda_graph_buffers = None + meta.num_sms = torch.cuda.get_device_properties(0).multi_processor_count + meta.runtime_tokens_per_gen_step = 0 + meta.create_expanded_buffers(capture_graph=False) + meta.prompt_lens_cpu = torch.zeros(max_seqs, dtype=torch.int32) + return meta + + +def _expand(meta, *, num_contexts, row_req_idx, num_seqs, block_offsets): + """Drive the private expansion with an explicit row->request map.""" + ngt = len(row_req_idx) + meta.num_contexts = num_contexts + meta._ragged_num_rows = ngt + meta.ragged_verify_lens = [1] * ngt + meta.kv_cache_block_offsets = block_offsets + meta.row_req_idx_host[:ngt].copy_( + torch.tensor(row_req_idx, dtype=torch.long)) + meta.row_kv_lens_host[:ngt].fill_(7) + meta.row_kv_correction_host[:ngt].zero_() + meta.prompt_lens_cpu[:num_seqs].fill_(1) + meta._prepare_token_major_gen_rows( + torch.full((num_seqs, ), 7, dtype=torch.int32)) + return meta._attn_num_rows + + +@pytest.mark.parametrize("num_pools", [1, 2]) +def test_row_block_table_matches_op_layout(num_pools): + """The expanded table keeps the op's rank, axis order and pool/KV axes.""" + torch.cuda.init() + max_blocks, max_seqs = 8, 4 + meta = _metadata(num_pools=num_pools, max_blocks=max_blocks, + max_seqs=max_seqs) + + src = torch.arange(num_pools * max_seqs * 2 * max_blocks, + dtype=torch.int32, + device="cuda").reshape(num_pools, max_seqs, 2, + max_blocks) + # Request 2 verifies three positions, request 0 one: a ragged batch, which + # is the whole point -- rows per request differ. + row_req_idx = [0, 2, 2, 2] + rows = _expand(meta, + num_contexts=0, + row_req_idx=row_req_idx, + num_seqs=max_seqs, + block_offsets=src) + assert rows == len(row_req_idx) + + got = meta.token_major_gen_view().kv_cache_block_offsets + assert got.dim() == src.dim(), ( + f"expanded table has rank {got.dim()} but the attention op reads rank " + f"{src.dim()}; substituting it wholesale would misread every axis") + assert got.shape[0] == num_pools and got.shape[2] == 2 + assert got.shape[1] == rows + # Every row must carry its own request's blocks, on every pool and both + # directions. + for row, req in enumerate(row_req_idx): + torch.testing.assert_close(got[:, row], src[:, req]) + + +def test_row_indices_stay_in_bounds_of_the_sequence_axis(): + """A request id must index num_seqs, never a shorter axis. + + This is the assertion the device-side gather was making. Checking it here + means the failure is a readable Python error rather than a CUDA abort that + takes the whole 8-rank job with it. + """ + torch.cuda.init() + meta = _metadata(num_pools=1, max_blocks=8, max_seqs=4) + + src = torch.zeros(1, 4, 2, 8, dtype=torch.int32, device="cuda") + rows = _expand(meta, + num_contexts=0, + row_req_idx=[0, 1, 3, 3], + num_seqs=4, + block_offsets=src) + idx = meta.attn_row_req_idx_cuda[:rows] + assert int(idx.min()) >= 0 + assert int(idx.max()) < src.shape[1], ( + f"row->request map reaches {int(idx.max())} but the sequence axis is " + f"{src.shape[1]} long") + + +def test_context_rows_precede_generation_rows(): + """Contexts keep row 0..nc-1 so the ops' generation offset still lands.""" + torch.cuda.init() + meta = _metadata(num_pools=1, max_blocks=8, max_seqs=6) + + src = torch.arange(1 * 6 * 2 * 8, dtype=torch.int32, + device="cuda").reshape(1, 6, 2, 8) + nc = 2 + rows = _expand(meta, + num_contexts=nc, + row_req_idx=[0, 1, 1], + num_seqs=6, + block_offsets=src) + assert rows == nc + 3 + idx = meta.attn_row_req_idx_host[:rows].tolist() + assert idx[:nc] == list(range(nc)), "context rows must be identity-mapped" + # Generation row->request ids are batch-relative and must be shifted past + # the contexts, or a generation row would read a context's blocks. + assert idx[nc:] == [nc + 0, nc + 1, nc + 1] + # host_context_lengths is the ops' shape oracle: one token per generation + # row, real length for contexts. + assert meta.attn_row_prompt_lens_cpu[nc:rows].tolist() == [1, 1, 1] + + +def test_every_runtime_view_shares_the_row_count(): + """All views the op reads must agree on batch_size, which is the row count. + + ``TrtllmAttention.forward`` asserts this directly, and the C++ op indexes + ``request_types`` over its whole ``num_seqs`` -- so a view left at the + request count is an out-of-bounds read there rather than a shape error. + """ + torch.cuda.init() + meta = _metadata(num_pools=1, max_blocks=8, max_seqs=6) + src = torch.zeros(1, 6, 2, 8, dtype=torch.int32, device="cuda") + nc = 1 + rows = _expand(meta, + num_contexts=nc, + row_req_idx=[0, 0, 1], + num_seqs=6, + block_offsets=src) + + view = meta.token_major_gen_view() + for name in ("sequence_length", "host_past_key_value_lengths", + "host_context_lengths", "prompt_lens_cuda", + "host_request_types"): + got = getattr(view, name).shape[0] + assert got == rows, ( + f"{name} has {got} entries but the ops will read {rows} rows") + assert view.kv_cache_block_offsets.shape[1] == rows + + # kCONTEXT is 0, kGENERATION is 1; the op checks every row at or past + # num_contexts is a generation row. + types = view.host_request_types.tolist() + assert types[:nc] == [0] * nc + assert types[nc:] == [1] * (rows - nc) diff --git a/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py b/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py index ff5743430938..f905b06f38ad 100644 --- a/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py +++ b/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py @@ -28,17 +28,25 @@ stall that delays the H2D), enqueues batch A's async H2D behind the stall, then lets the host immediately issue batch B before A's copy can drain. With the bug present, batch A's device tensor ends up holding batch B's offsets. + +``KVCacheManagerV2`` shared the bug through the persistent +``host_kv_cache_block_offsets`` it fed to the same async H2D — worse, in V2 +that buffer's rows are the live page-index buffers the C++ allocator rewrites +in place at scheduling time. The V2 twin at the bottom drives the equivalent +window against its snapshot fix (``_stage_host_block_offsets_for_copy``). """ import pytest import torch +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings import DataType, LayerType from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp from tensorrt_llm.bindings.internal.batch_manager import CacheType from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2._utils import init_cuda_once # Small geometry so each sequence spans several blocks and the two batches get # visibly different physical block indices. @@ -52,6 +60,10 @@ # Long GPU stall so the host reliably wins the race and overwrites the staging # buffer before batch A's H2D drains. Generous to stay deterministic on fast GPUs. _SLEEP_CYCLES = 2_000_000_000 +# Overwrites every V2 page index during the race window. Any value distinct +# from real page indices (and from BAD_PAGE_INDEX == -1) works; the copies +# only move it as data, never index with it. +_V2_CLOBBER_PAGE_INDEX = 0x0BAD def _build_manager(): @@ -89,8 +101,13 @@ def _build_manager(): def _empty_device_offsets(mgr): + # Dim 0 mirrors production's kv_cache_block_offsets allocation (see + # TrtllmAttentionMetadata): V2 sizes it by num_attention_op_pools, which + # equals num_pools unless SWA scratch reuse widens it to one virtual pool + # per layer. V1 has no such attribute. + num_pools = getattr(mgr, "num_attention_op_pools", mgr.num_pools) return torch.zeros( - mgr.num_pools, 2 * _BATCH, 2, mgr.max_blocks_per_seq, dtype=torch.int32, device="cuda" + num_pools, 2 * _BATCH, 2, mgr.max_blocks_per_seq, dtype=torch.int32, device="cuda" ) @@ -186,3 +203,99 @@ def test_copy_batch_block_offsets_survives_overlap_overwrite(): ) # And batch B's copy is independently correct. assert torch.equal(dst_b[:, :_BATCH], ref_b) + + +def _build_manager_v2(enable_swa_scratch_reuse): + """Plain full-attention V2 manager with the same small geometry as V1's. + + ``enable_swa_scratch_reuse=True`` routes ``copy_batch_block_offsets`` + through ``_copy_batch_block_offsets_per_layer`` — the branch production + DSV4 serving takes and the one the snapshot fix primarily targets. + ``False`` covers the flat ``copy_batch_block_offsets_to_device`` branch, + which stages through the same helper. + """ + init_cuda_once() + kv_cache_config = KvCacheConfig( + max_tokens=4096, + enable_block_reuse=False, + enable_swa_scratch_reuse=enable_swa_scratch_reuse, + ) + return KVCacheManagerV2( + kv_cache_config, + CacheType.SELF, + num_layers=_NUM_LAYERS, + num_kv_heads=_NUM_KV_HEADS, + head_dim=_HEAD_DIM, + tokens_per_block=_TOKENS_PER_BLOCK, + max_seq_len=_MAX_SEQ_LEN, + max_batch_size=2 * _BATCH, + mapping=Mapping(world_size=1, tp_size=1, pp_size=1), + dtype=DataType.HALF, + vocab_size=32000, + # Keep the pool solver's typical step at test scale instead of the + # constructor's 8192-token default. + max_num_tokens=_MAX_SEQ_LEN, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") +@pytest.mark.parametrize("swa_scratch_reuse", [True, False], ids=["per_layer_scratch", "flat"]) +def test_kv_cache_manager_v2_copy_batch_block_offsets_survives_host_rewrite(swa_scratch_reuse): + """V2 twin of the overlap-overwrite test above. + + Batch A's copy is enqueued behind a stream stall, then the persistent + ``host_kv_cache_block_offsets`` is rewritten in place before the copy can + drain. The rewrite is direct rather than a free-A/allocate-B cycle: the + C++ allocator mutates exactly these rows in place (through the + memoryviews registered by ``_create_kv_cache`` / + ``_restore_page_index_bufs``), but a re-allocation may hand the freed + pages back in the same order and reproduce identical rows, masking the + race. Batch A's device tensor must keep batch A's offsets. + """ + mgr = _build_manager_v2(swa_scratch_reuse) + assert mgr.enable_swa_scratch_reuse == swa_scratch_reuse + requests = None + saved_host_offsets = None + try: + ids = list(range(1, 1 + _BATCH)) + toks = [_TOKENS_PER_SEQ] * _BATCH + requests = mgr.add_dummy_requests(request_ids=ids, token_nums=toks, prepare_resource=True) + assert requests is not None, "test setup invalid: dummy request allocation failed" + + # Also warms the (maybe-compiled) per-layer copy op outside the race + # window, so the stalled call below enqueues without host-side + # compilation pauses that would let the H2D drain early. + ref = _reference_offsets(mgr, ids) + saved_host_offsets = mgr.host_kv_cache_block_offsets.clone() + + dst = _empty_device_offsets(mgr) + torch.cuda.synchronize() + + # Drive the scheduling window: stall the stream, enqueue batch A's + # async H2D behind the stall, then rewrite the persistent host rows + # before the copy can drain. + torch.cuda._sleep(_SLEEP_CYCLES) + mgr.copy_batch_block_offsets(dst, ids, 1, _BATCH, _BATCH) + mgr.host_kv_cache_block_offsets.fill_(_V2_CLOBBER_PAGE_INDEX) + torch.cuda.synchronize() + + # Setup validity: a copy issued against the clobbered buffer must + # produce something else, or the rewrite wasn't observable. + clobbered = _reference_offsets(mgr, ids) + assert not torch.equal(clobbered, ref), ( + "test setup invalid: host-buffer rewrite did not change the copied offsets" + ) + + assert torch.equal(dst[:, :_BATCH], ref), ( + "nvbug 6293536 (V2): batch A's async H2D read the rewritten " + "persistent host_kv_cache_block_offsets instead of a private " + "snapshot (_stage_host_block_offsets_for_copy)" + ) + finally: + if saved_host_offsets is not None: + # The rows are live page-index buffers; put them back before the + # allocator touches them again in free/shutdown. + mgr.host_kv_cache_block_offsets.copy_(saved_host_offsets) + for req in requests or []: + mgr.free_resources(req) + mgr.shutdown() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_accept_stats.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_accept_stats.py new file mode 100644 index 000000000000..080294abb8a4 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_accept_stats.py @@ -0,0 +1,363 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark accept-caps and ragged-verify observability (hardware-agnostic, CPU). + +Covers `apply_accept_caps` clamping semantics, the cap-accept diagnostic mode's +counters (`cap_trim_tokens`, trim shape), `trim_ratio` row-base accounting under +CUDA-graph padding, and trim-regret acceptance stats. +""" + +import os +from unittest.mock import patch + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_observability import ( + DSparkRaggedStats, RaggedVerifyMode) +from tensorrt_llm._torch.speculative.interface import apply_accept_caps + +MAX_DRAFT_LEN = 5 + + +def _cap_stats(mode=RaggedVerifyMode.CAP_ACCEPT): + return DSparkRaggedStats(mode=mode, max_draft_len=MAX_DRAFT_LEN) + + +def _compact_stats(max_draft_len=5): + return DSparkRaggedStats(mode=RaggedVerifyMode.COMPACT, + max_draft_len=max_draft_len) + + +class _Meta: + """Just the two fields `apply_accept_caps` reads.""" + + def __init__(self, caps=None, track_trim=True, rows=8): + self.accept_caps = (None if caps is None else torch.tensor( + caps, dtype=torch.int32)) + # Persistent per-request buffer, sized like the engine's. + self.accept_cap_trim = (torch.zeros(rows, dtype=torch.int32) + if track_trim else None) + + +# --- the cap itself ---------------------------------------------------------- +# Units: `num_accepted_tokens` is `accepted drafts + 1`, so a window of w +# drafted positions is a cap of `w + 1`, matching `verify_lens`. + + +def test_no_caps_is_a_no_op(): + """Every other mode goes through this call.""" + counts = torch.tensor([3, 1, 4], dtype=torch.int32) + apply_accept_caps(counts, 0, _Meta(caps=None)) + assert counts.tolist() == [3, 1, 4] + + +def test_counts_beyond_their_windows_are_clamped_and_the_loss_accumulated(): + """The trim buffer holds this step's loss, not a running sum: repeating the + same step leaves the same losses.""" + meta = _Meta(caps=[3, 5, 6]) + for _ in range(3): + counts = torch.tensor([6, 5, 2], dtype=torch.int32) + apply_accept_caps(counts, 0, meta) + # Request 2's cap (6) is wider than its count (2): never raised. + assert counts.tolist() == [3, 5, 2] + # Only the first request lost anything: 6 -> 3. And we can say WHICH. + assert meta.accept_cap_trim[:3].tolist() == [3, 0, 0] + + +def test_context_requests_are_left_alone(): + """Caps are per generation request; contexts sit ahead of them.""" + counts = torch.tensor([9, 9, 6, 6], dtype=torch.int32) + meta = _Meta(caps=[2, 3]) + apply_accept_caps(counts, 2, meta) + assert counts.tolist() == [9, 9, 2, 3], "a context count was clamped" + # Context slots are zeroed, not left stale. + assert meta.accept_cap_trim[:4].tolist() == [0, 0, 4, 3] + + # Degenerate endpoint: a generation-free batch must not crash or clamp. + counts = torch.tensor([4, 4], dtype=torch.int32) + apply_accept_caps(counts, 2, _Meta(caps=[2, 2])) + assert counts.tolist() == [4, 4] + + +def test_missing_accumulator_still_caps(): + """The clamp is correctness; the counter is only measurement.""" + counts = torch.tensor([6], dtype=torch.int32) + meta = _Meta(caps=[2], track_trim=False) + apply_accept_caps(counts, 0, meta) + assert counts.tolist() == [2] + + +@pytest.mark.parametrize("accepted", range(0, MAX_DRAFT_LEN + 2)) +@pytest.mark.parametrize("window", range(0, MAX_DRAFT_LEN + 1)) +def test_nothing_is_lost_or_invented(accepted, window): + """committed + discarded accounts for every position the target produced.""" + counts = torch.tensor([accepted + 1], dtype=torch.int32) + meta = _Meta(caps=[window + 1]) + apply_accept_caps(counts, 0, meta) + committed = int(counts[0]) + trim = int(meta.accept_cap_trim[0]) + assert committed + trim == accepted + 1 + assert committed >= 1, "a step that commits nothing cannot make progress" + assert committed <= window + 1 + + +# --- what the counters must say ---------------------------------------------- + + +def test_cap_trim_is_recorded_per_request(): + stats = _cap_stats() + # Given 5, allowed 2 -> 3 accepted positions thrown away. + stats.record_acceptance(accepted=2, window=2, cap_trim=3) + # Died on its own inside a wide window: the schedule cost it nothing. + stats.record_acceptance(accepted=1, window=4, cap_trim=0) + + assert stats.cap_trim_tokens == 3 + assert stats.requests_cap_trimmed == 1 + assert stats.accept_len == pytest.approx(1.5) + assert stats.accept_loss_per_request == pytest.approx(1.5) + # accept_len + accept_loss is what a no-trim run would have accepted. + assert (stats.accept_len + + stats.accept_loss_per_request) == pytest.approx(3.0) + + # And the summary carries the measurement. + summary = stats.summary() + assert summary["mode"] == "cap-accept" + assert summary["cap_trim_tokens"] == 3 + assert summary["accept_loss_per_request"] == pytest.approx(1.5) + for key in ("requests_cap_trimmed", "cap_trim_max", "cap_trim_hist", + "cap_trim_concentration"): + assert key in summary, f"{key} missing from the summary" + + +def test_the_same_total_loss_is_distinguished_by_its_shape(): + """Equal totals with different per-request shapes must stay + distinguishable: spread loss vs. loss concentrated in one request.""" + spread = _cap_stats() + for _ in range(10): + spread.record_acceptance(accepted=1, window=1, cap_trim=2) + + concentrated = _cap_stats() + concentrated.record_acceptance(accepted=1, window=1, cap_trim=20) + for _ in range(9): + concentrated.record_acceptance(accepted=3, window=5, cap_trim=0) + + assert spread.cap_trim_tokens == concentrated.cap_trim_tokens == 20 + assert spread.accept_loss_per_request == pytest.approx( + concentrated.accept_loss_per_request) + + # ...and the shape separates them. + assert spread.cap_trim_concentration == pytest.approx(1.0) + assert concentrated.cap_trim_concentration == pytest.approx(0.1) + assert spread.cap_trim_max == 2 + assert concentrated.cap_trim_max == 20 + assert spread.cap_trim_hist == {2: 10} + assert concentrated.cap_trim_hist == {20: 1} + + +def test_the_loss_is_absent_rather_than_zero_when_not_measured(): + """`compact` cannot produce this number, and must not appear to.""" + stats = _cap_stats(mode=RaggedVerifyMode.COMPACT) + stats.record_acceptance(accepted=3, window=3) + assert stats.cap_trim_tokens == 0 + assert stats.cap_trim_concentration == 0.0 + assert stats.cap_trim_max == 0 + # The bound that IS available under compact still works. + assert stats.trimmed_hit_ceiling == 1 + + +def test_cap_accept_must_not_be_credited_with_a_compute_saving(): + """cap-accept submits the full block, so delivered tokens must not be + derived from the windows and `trim_ratio` must stay 0.""" + stats = _cap_stats() + num_reqs = 4 + full_block = num_reqs * (1 + MAX_DRAFT_LEN) + stats.record_step(num_gen_requests=num_reqs, + verify_lens=[5, 3, 2, 1], + bucket=None, + delivered=full_block) + + assert stats.delivered_tokens == full_block + assert stats.ceiling_tokens == full_block + assert stats.trim_ratio == 0.0, "cap-accept saves no compute" + # The schedule is still visible -- that is the whole point of the mode. + assert stats.window_tokens == sum(1 + v for v in (5, 3, 2, 1)) + assert stats.steps_ragged == 1 + + +def test_compact_still_derives_delivered_from_the_bucket(): + """The default path is unchanged by the new argument.""" + stats = _cap_stats(mode=RaggedVerifyMode.COMPACT) + stats.record_step(num_gen_requests=4, + verify_lens=[5, 3, 2, 1], + bucket=16) + assert stats.delivered_tokens == 16 + assert stats.trim_ratio > 0.0 + + +# --- mode plumbing ----------------------------------------------------------- + + +def test_cap_accept_computes_windows_but_does_not_trim_the_token_axis(): + """The two predicates that route the whole feature.""" + assert RaggedVerifyMode.CAP_ACCEPT.computes_windows + assert not RaggedVerifyMode.CAP_ACCEPT.trims_submitted_tokens + + assert RaggedVerifyMode.COMPACT.computes_windows + assert RaggedVerifyMode.COMPACT.trims_submitted_tokens + + assert not RaggedVerifyMode.STATIC.computes_windows + assert not RaggedVerifyMode.STATIC.trims_submitted_tokens + + +def test_a_slot_left_unwritten_cannot_inherit_the_previous_occupant_s_loss(): + """cap_trim_lens is persistent and slot-indexed: a step producing no caps + must still zero this batch's slots or recycled slots report stale loss.""" + import inspect + + from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSamplerBase + + src = inspect.getsource(SpecSamplerBase.sample_async) + assert "zeros_like" in src and "cap_trim_lens" in src, ( + "_process_outputs no longer writes zeros when the step produced no " + "cap_trim_lens; recycled slots will report a stale request's loss") + + +def test_the_cuda_graph_padding_dummy_carries_a_cap(): + """The CUDA-graph padding dummy must carry a verify cap, or cap-accept + silently degrades to static on every padded step.""" + import inspect + + from tensorrt_llm._torch.pyexecutor import cuda_graph_runner + + src = inspect.getsource(cuda_graph_runner.CUDAGraphRunner._get_padded_batch) + assert "py_verify_cap" in src, ( + "the padded-batch builder no longer stamps the dummy with a verify " + "cap; cap-accept will silently degrade to static on every padded step") + # And it must be cleared, not left stale: the dummy object is cached + # across steps and shared between them. + assert "if cap_accept else None" in src, ( + "the dummy's cap is not cleared for non-cap-accept batches, so a " + "stale cap would trim a static step") + + +def test_the_capture_set_follows_the_mode_not_the_config_flag(): + """Batch shaping must ask whether the token axis actually shrinks, not + whether ragged was requested: cap-accept always submits the full block.""" + from types import SimpleNamespace + + from tensorrt_llm._torch.speculative.dspark_observability import ( + resolve_ragged_verify_mode, trims_submitted_tokens) + + ragged_cfg = SimpleNamespace(enable_ragged_verify=True) + + with patch.dict(os.environ, + {"TLLM_DSPARK_RAGGED_VERIFY_MODE": "cap-accept"}): + assert resolve_ragged_verify_mode( + ragged_cfg) is RaggedVerifyMode.CAP_ACCEPT + assert not trims_submitted_tokens(ragged_cfg), ( + "cap-accept submits the full block; a capture set built for it " + "must be the uniform one") + + with patch.dict(os.environ, {"TLLM_DSPARK_RAGGED_VERIFY_MODE": "compact"}): + assert trims_submitted_tokens(ragged_cfg) + + # Unset: the config flag decides, which is the pre-existing behaviour. + with patch.dict(os.environ, {}, clear=True): + assert trims_submitted_tokens(ragged_cfg) + assert not trims_submitted_tokens( + SimpleNamespace(enable_ragged_verify=False)) + assert not trims_submitted_tokens(None) + + +def test_accept_caps_is_not_the_ragged_flag(): + """`accept_caps` must not imply ragged buffer packing: cap-accept's draft + buffers are the ordinary rectangle.""" + from tensorrt_llm._torch.speculative.interface import SpecMetadata + + fields = SpecMetadata.__dataclass_fields__ + assert "accept_caps" in fields + assert "accept_cap_trim" in fields + + meta = _Meta(caps=[2, 3]) + # The mock stands in for a metadata object carrying caps but no windows. + assert getattr(meta, "verify_lens", None) is None + + +# --- trim_ratio row-base under CUDA-graph padding ---------------------------- + + +def test_padding_rows_do_not_make_the_trim_look_negative(): + """A top-tier step on padded rows saved nothing -- and lost nothing.""" + stats = _compact_stats(max_draft_len=5) + # 3 real requests, graph padded to 4 rows, every window at the full block. + # delivered = 4 * 6 = 24. Against a ceiling of 3 * 6 = 18 that is -0.333. + stats.record_step(num_gen_requests=3, verify_lens=[5, 5, 5], + bucket=4 * 6, padded_bs=4) + summary = stats.summary() + assert summary["trim_ratio"] == pytest.approx(0.0), ( + f"a step that verified the full block on every row trimmed nothing, so " + f"the ratio is 0 -- got {summary['trim_ratio']}, which came from " + f"comparing {summary['delivered_tokens']} padded-row tokens against a " + f"ceiling counted over {3} real requests") + # The padded row count itself must land in the summary: the executor once + # passed padded_bs=None unconditionally, and an always-empty padded_bs_hist + # is indistinguishable from "padding never happened". + assert summary["padded_bs_hist"] == {4: 1} + + +def test_a_real_trim_still_reads_as_a_saving(): + """The fix must not flatten genuine trimming to zero.""" + stats = _compact_stats(max_draft_len=5) + # 4 rows, no padding, every window trimmed to 2 -> bucket 4*3 = 12 against a + # ceiling of 4*6 = 24. + stats.record_step(num_gen_requests=4, verify_lens=[2, 2, 2, 2], + bucket=4 * 3, padded_bs=4) + assert stats.summary()["trim_ratio"] == pytest.approx(0.5) + + +def test_paths_without_a_bucket_keep_the_local_row_base(): + """Only the bucket path rebases to the padded row count; widening the + no-window and cap-accept paths would invent a saving out of padding.""" + no_window = _compact_stats(max_draft_len=5) + no_window.record_step(num_gen_requests=3, verify_lens=None, padded_bs=8) + assert no_window.summary()["trim_ratio"] == pytest.approx(0.0) + + explicit = _compact_stats(max_draft_len=5) + # cap-accept: windows are computed but the full block is submitted anyway. + explicit.record_step(num_gen_requests=3, verify_lens=[2, 3, 1], + bucket=8 * 6, padded_bs=8, delivered=3 * 6) + assert explicit.summary()["trim_ratio"] == pytest.approx(0.0) + + +# --- trim regret ------------------------------------------------------------- + + +def test_trim_regret_counts_only_drafts_alive_at_the_cut(): + """Regret separates "trimming was free, those drafts would have died + anyway" from "trimming discarded acceptance".""" + stats = _compact_stats() + # Given the full block and died early: not trimmed, no regret. + stats.record_acceptance(accepted=2, window=5) + # Trimmed to 3 and died at 1: the cut cost nothing. + stats.record_acceptance(accepted=1, window=3) + # Trimmed to 3 and accepted all 3: still alive at the cut -> regret. + stats.record_acceptance(accepted=3, window=3) + + assert stats.requests_scored == 3 + assert stats.requests_trimmed == 2 + assert stats.trimmed_hit_ceiling == 1 + assert stats.trim_regret_rate == 0.5 + assert stats.accept_len == pytest.approx(2.0) + summary = stats.summary() + for key in ("accept_len", "requests_scored", "requests_trimmed", + "trim_regret_rate"): + assert key in summary, f"{key} missing from the summary" + + # Nothing trimmed: the regret rate is 0/0-safe and stays 0. + untrimmed = _compact_stats() + for accepted in (0, 3, 5): + untrimmed.record_acceptance(accepted=accepted, window=5) + assert untrimmed.requests_trimmed == 0 + assert untrimmed.trim_regret_rate == 0.0 + assert untrimmed.summary()["accept_len"] == pytest.approx(8 / 3, abs=1e-4) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py deleted file mode 100644 index 6ed4f1cba050..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ /dev/null @@ -1,454 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Unit tests for the DSpark captured-context attention primitives (CPU). - -The two primitives are validated against fully independent computations: -``dspark_sparse_attn`` vs (a) a per-element loop reference of the kernel formula -and (b) ``torch.nn.functional.scaled_dot_product_attention`` for the no-sink, -all-valid case; ``get_dspark_topk_idxs`` vs the reference index formula. -""" - -import types -from unittest.mock import Mock - -import pytest -import torch -import torch.nn.functional as F - -import tensorrt_llm._torch.models.modeling_dspark as modeling_dspark -from tensorrt_llm._torch.models.dspark.attention import ( - apply_dspark_rotary, - dspark_attention_forward, - dspark_sparse_attn, - get_dspark_topk_idxs, - precompute_dspark_freqs_cis, -) -from tensorrt_llm._torch.models.modeling_dspark import DSparkDraftModel - - -def test_rope_table_is_cached_once_per_device(): - model = types.SimpleNamespace( - _attn_params={"rope_head_dim": 16}, - _freqs_cap=64, - _rope_theta=10000.0, - _freqs_table_cache={}, - ) - - first = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) - second = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) - - assert first.data_ptr() == second.data_ptr() - assert len(model._freqs_table_cache) == 1 - positions = torch.tensor([1, 17, 63]) - expected = precompute_dspark_freqs_cis(16, 64, rope_theta=10000.0) - torch.testing.assert_close(first[positions], expected[positions]) - - -def test_dspark_block_uses_stage_id_as_attention_layer_idx(monkeypatch): - captured = {} - - def fake_decoder_layer_init( - self, - model_config, - layer_idx, - aux_stream_dict, - attention_layer_idx=None, - mapping_with_cp=None, - disable_post_moe_fusion=False, - ): - torch.nn.Module.__init__(self) - self.model_config = model_config - self.config = model_config.pretrained_config - self.layer_idx = layer_idx - captured.update( - layer_idx=layer_idx, - attention_layer_idx=attention_layer_idx, - aux_stream_dict=aux_stream_dict, - mapping_with_cp=mapping_with_cp, - disable_post_moe_fusion=disable_post_moe_fusion, - ) - - monkeypatch.setattr( - modeling_dspark.DeepseekV4DecoderLayer, - "__init__", - fake_decoder_layer_init, - ) - model_config = types.SimpleNamespace( - pretrained_config=types.SimpleNamespace(vocab_size=128, hc_mult=2), - spec_config=None, - ) - - block = modeling_dspark.DSparkBlock( - model_config, - layer_idx=10, - aux_stream_dict={}, - stage_id=1, - num_stages=3, - num_capture_layers=0, - ) - - assert block.layer_idx == captured["layer_idx"] == 10 - assert captured["attention_layer_idx"] == block.stage_id == 1 - assert captured["disable_post_moe_fusion"] is True - - -@pytest.mark.parametrize("enable_fused_hc", [True, False]) -def test_forward_stage_honors_enable_fused_hc(monkeypatch, enable_fused_hc): - """The draft stage must use the inherited fused-HC rollback setting.""" - torch.manual_seed(71) - num_requests, block_size, hc_mult, hidden_size = 1, 2, 2, 3 - h = torch.randn(num_requests, block_size, hc_mult, hidden_size) - attention_input = torch.randn(num_requests, block_size, hidden_size) - attention_output = torch.randn_like(attention_input) - mid_residual = torch.randn_like(h) - attention_post_mix = torch.randn(num_requests, block_size, hc_mult, 1) - attention_comb_mix = torch.randn(num_requests, block_size, hc_mult, hc_mult) - ffn_post_mix = torch.randn_like(attention_post_mix) - ffn_comb_mix = torch.randn_like(attention_comb_mix) - raw_ffn_input = torch.randn_like(attention_input) - normed_ffn_input = torch.randn_like(attention_input) - moe_output = torch.randn(num_requests * block_size, hidden_size) - final_h = torch.randn_like(h) - events = [] - - def record(name, result): - def call(*args, **kwargs): - events.append(name) - return result - - return call - - monkeypatch.setattr( - modeling_dspark, - "dspark_attention_forward", - Mock(return_value=attention_output), - ) - - hc_attn = types.SimpleNamespace( - pre_mapping=Mock(return_value=(attention_post_mix, attention_comb_mix, attention_input)), - post_mapping=Mock(side_effect=record("attention_post", mid_residual)), - ) - hc_ffn = types.SimpleNamespace( - fused_hc=Mock( - side_effect=record( - "fused", - (mid_residual, ffn_post_mix, ffn_comb_mix, normed_ffn_input), - ) - ), - pre_mapping=Mock( - side_effect=record("ffn_pre", (ffn_post_mix, ffn_comb_mix, raw_ffn_input)) - ), - post_mapping=Mock(side_effect=record("ffn_post", final_h)), - ) - post_attention_layernorm = Mock(side_effect=record("ffn_norm", normed_ffn_input)) - post_attention_layernorm.weight = torch.ones(hidden_size) - post_attention_layernorm.variance_epsilon = 1e-6 - stage = types.SimpleNamespace( - enable_fused_hc=enable_fused_hc, - hc_attn=hc_attn, - hc_ffn=hc_ffn, - input_layernorm=Mock(side_effect=lambda tensor: tensor), - post_attention_layernorm=post_attention_layernorm, - mlp=Mock(return_value=moe_output), - _dspark_attn={}, - ) - model = types.SimpleNamespace( - use_real_mla=False, - _attn_params={"window_size": 2, "head_dim": 1}, - model_config=types.SimpleNamespace( - mapping=types.SimpleNamespace(enable_attention_dp=False, tp_size=8) - ), - ) - - actual = DSparkDraftModel._forward_stage( - model, - stage, - h, - torch.randn(num_requests, hidden_size), - 1, - torch.empty(0), - torch.zeros(num_requests, block_size, dtype=torch.long), - ) - - assert actual is final_h - torch.testing.assert_close( - stage.mlp.call_args.args[0], - normed_ffn_input.reshape(num_requests * block_size, hidden_size), - ) - # Non-attention-DP multi-GPU (enable_attention_dp=False, tp_size>1): the draft - # MoE must all-reduce its TP-sharded output, mirroring the target MoE. The - # attention-DP and single-GPU paths keep it disabled. - assert stage.mlp.call_args.kwargs["final_all_reduce_params"].enable_allreduce is True - if enable_fused_hc: - assert events == ["fused", "ffn_post"] - hc_ffn.fused_hc.assert_called_once() - hc_attn.post_mapping.assert_not_called() - hc_ffn.pre_mapping.assert_not_called() - post_attention_layernorm.assert_not_called() - fused_kwargs = hc_ffn.fused_hc.call_args.kwargs - assert fused_kwargs["norm_weight"] is post_attention_layernorm.weight - assert fused_kwargs["norm_eps"] == post_attention_layernorm.variance_epsilon - else: - assert events == ["attention_post", "ffn_pre", "ffn_norm", "ffn_post"] - hc_ffn.fused_hc.assert_not_called() - hc_attn.post_mapping.assert_called_once() - hc_ffn.pre_mapping.assert_called_once_with(mid_residual) - post_attention_layernorm.assert_called_once_with(raw_ffn_input) - - -def _ref_precompute_freqs_cis(dim, seqlen, base): - """DeepSpec precompute_freqs_cis with original_seq_len == 0 (no YaRN).""" - freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - t = torch.arange(seqlen) - freqs = torch.outer(t, freqs) - return torch.polar(torch.ones_like(freqs), freqs) - - -def _ref_apply_rotary_emb(x, freqs_cis, inverse=False): - """DeepSpec apply_rotary_emb (returns a fresh tensor instead of in-place).""" - xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) - if inverse: - freqs_cis = freqs_cis.conj() - if xc.ndim == 3: - fc = freqs_cis.view(1, xc.size(1), xc.size(-1)) - else: - fc = freqs_cis.view(1, xc.size(1), 1, xc.size(-1)) - return torch.view_as_real(xc * fc).flatten(-2).to(x.dtype) - - -def _loop_reference(q, kv, attn_sink, topk_idxs, scale): - """Obvious, slow per-(b,m,h) reference of the exact kernel math.""" - b, m, h, d = q.shape - out = torch.zeros(b, m, h, d, dtype=torch.float32) - qf, kvf, sink = q.float(), kv.float(), attn_sink.float() - for bi in range(b): - for mi in range(m): - idxs = topk_idxs[bi, mi].tolist() - for hi in range(h): - scores, vecs = [], [] - for j in idxs: - if j < 0: - continue - k = kvf[bi, j] - scores.append(torch.dot(qf[bi, mi, hi], k) * scale) - vecs.append(k) - if not scores: - continue - s = torch.stack(scores) - smax = s.max() - p = torch.exp(s - smax) - denom = p.sum() + torch.exp(sink[hi] - smax) - num = (p.unsqueeze(-1) * torch.stack(vecs)).sum(0) - out[bi, mi, hi] = num / denom - return out - - -@pytest.mark.parametrize("seed", [0, 1, 2]) -def test_sparse_attn_matches_loop_reference(seed): - torch.manual_seed(seed) - b, m, h, d, n, topk = 2, 5, 3, 16, 40, 12 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, n, d) - attn_sink = torch.randn(h) - idx = torch.stack( - [torch.stack([torch.randperm(n)[:topk] for _ in range(m)]) for _ in range(b)] - ).int() - scale = d**-0.5 - got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() - ref = _loop_reference(q, kv, attn_sink, idx, scale) - torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) - - -def test_sparse_attn_no_sink_matches_sdpa(): - """With a -inf sink (no sink mass) and all-valid contiguous indices, the - primitive must equal standard scaled-dot-product attention over the gathered - KV — an independent implementation.""" - torch.manual_seed(0) - b, m, h, d, topk = 2, 4, 5, 16, 9 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, topk, d) # n == topk, attend to all - attn_sink = torch.full((h,), float("-inf")) - idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() - scale = d**-0.5 - got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() - - # SDPA: q [b,h,m,d], k/v [b,h,topk,d] (broadcast the shared KV over heads). - qh = q.permute(0, 2, 1, 3) - kvh = kv.unsqueeze(1).expand(b, h, topk, d) - ref = F.scaled_dot_product_attention(qh, kvh, kvh, scale=scale).permute(0, 2, 1, 3) - torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) - - -def test_sparse_attn_sink_reduces_mass(): - """A finite sink must strictly shrink the attention output magnitude vs an - infinitely-negative (disabled) sink, because it adds denominator mass only.""" - torch.manual_seed(0) - b, m, h, d, topk = 1, 2, 2, 16, 6 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, topk, d) - idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() - scale = d**-0.5 - no_sink = dspark_sparse_attn(q, kv, torch.full((h,), float("-inf")), idx, scale) - with_sink = dspark_sparse_attn(q, kv, torch.zeros(h), idx, scale) - assert with_sink.abs().sum() < no_sink.abs().sum() - - -def test_sparse_attn_masked_indices_excluded(): - """An index of -1 must be excluded exactly (equiv. to dropping that column).""" - torch.manual_seed(0) - b, m, h, d = 1, 1, 2, 16 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, 5, d) - sink = torch.full((h,), float("-inf")) - scale = d**-0.5 - full = torch.tensor([[[0, 1, 2, 3]]]).int() - masked = torch.tensor([[[0, 1, 2, -1]]]).int() - drop3 = torch.tensor([[[0, 1, 2]]]).int() - got_masked = dspark_sparse_attn(q, kv, sink, masked, scale) - got_drop = dspark_sparse_attn(q, kv, sink, drop3, scale) - torch.testing.assert_close(got_masked, got_drop, rtol=1e-5, atol=1e-5) - # And masking genuinely changes the result vs attending to position 3. - got_full = dspark_sparse_attn(q, kv, sink, full, scale) - assert not torch.allclose(got_full, got_masked, rtol=1e-3, atol=1e-3) - - -@pytest.mark.parametrize( - "start_pos,window,block", [(1, 128, 5), (3, 128, 5), (10, 4, 5), (200, 128, 6)] -) -def test_get_dspark_topk_idxs_matches_reference(start_pos, window, block): - bsz = 3 - got = get_dspark_topk_idxs(window, bsz, block, start_pos) - # Reference formula (DeepSpec get_dspark_topk_idxs). - ctx = torch.arange(min(window, start_pos + 1)) - blk = window + torch.arange(block) - ref_row = torch.cat([ctx, blk]).int() - assert got.shape == (bsz, block, ref_row.numel()) - for bi in range(bsz): - for mi in range(block): - torch.testing.assert_close(got[bi, mi], ref_row) - - -def test_get_dspark_topk_idxs_requires_generation(): - with pytest.raises(AssertionError): - get_dspark_topk_idxs(128, 1, 5, 0) - - -@pytest.mark.parametrize("rope_head_dim,seqlen", [(64, 16), (64, 1), (128, 8)]) -def test_precompute_freqs_cis_matches_reference(rope_head_dim, seqlen): - got = precompute_dspark_freqs_cis(rope_head_dim, seqlen, rope_theta=10000.0) - ref = _ref_precompute_freqs_cis(rope_head_dim, seqlen, 10000.0) - torch.testing.assert_close(got, ref) - - -@pytest.mark.parametrize("ndim", [3, 4]) -def test_apply_rotary_matches_reference(ndim): - torch.manual_seed(0) - b, s, h, rd = 2, 5, 4, 64 - x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) - fc = precompute_dspark_freqs_cis(rd, s) - got = apply_dspark_rotary(x, fc) - ref = _ref_apply_rotary_emb(x, fc) - torch.testing.assert_close(got, ref) - - -@pytest.mark.parametrize("ndim", [3, 4]) -def test_apply_rotary_inverse_roundtrip(ndim): - """De-rotation (inverse) must undo the forward rotation (property test).""" - torch.manual_seed(1) - b, s, h, rd = 2, 6, 3, 64 - x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) - fc = precompute_dspark_freqs_cis(rd, s) - roundtrip = apply_dspark_rotary(apply_dspark_rotary(x, fc), fc, inverse=True) - torch.testing.assert_close(roundtrip, x, rtol=1e-5, atol=1e-5) - - -def _make_attn_inputs(seed=0): - """Small synthetic DSpark attention inputs/weights (CPU bf16).""" - torch.manual_seed(seed) - dim, n_heads, head_dim, rd = 12, 4, 8, 4 - q_lora, o_lora, n_groups = 6, 5, 2 - window, block, start_pos = 8, 3, 5 - b = 2 - g = dict( - dim=dim, - n_heads=n_heads, - head_dim=head_dim, - rope_head_dim=rd, - q_lora=q_lora, - o_lora=o_lora, - n_groups=n_groups, - window=window, - block=block, - start_pos=start_pos, - b=b, - eps=1e-6, - softmax_scale=head_dim**-0.5, - ) - bf = torch.bfloat16 - g["x"] = torch.randn(b, block, dim, dtype=bf) - g["main_x"] = torch.randn(b, 1, dim, dtype=bf) - g["kv_cache0"] = torch.randn(b, window, head_dim, dtype=bf) - g["wq_a"] = torch.randn(q_lora, dim, dtype=bf) * 0.1 - g["wq_b"] = torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1 - g["wkv"] = torch.randn(head_dim, dim, dtype=bf) * 0.1 - g["wo_a"] = torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1 - g["wo_b"] = torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1 - g["q_norm"] = torch.ones(q_lora) - g["kv_norm"] = torch.ones(head_dim) - g["attn_sink"] = torch.randn(n_heads) - g["freqs"] = precompute_dspark_freqs_cis(rd, start_pos + 1 + block + 2) - return g - - -def _run(g): - return dspark_attention_forward( - g["x"], - g["main_x"], - g["start_pos"], - g["kv_cache0"], - wq_a=g["wq_a"], - q_norm_w=g["q_norm"], - wq_b=g["wq_b"], - wkv=g["wkv"], - kv_norm_w=g["kv_norm"], - wo_a=g["wo_a"], - wo_b=g["wo_b"], - attn_sink=g["attn_sink"], - n_heads=g["n_heads"], - head_dim=g["head_dim"], - rope_head_dim=g["rope_head_dim"], - n_groups=g["n_groups"], - o_lora_rank=g["o_lora"], - window_size=g["window"], - eps=g["eps"], - softmax_scale=g["softmax_scale"], - freqs_cis=g["freqs"], - ) - - -def test_attention_forward_shape_and_determinism(): - g = _make_attn_inputs() - o = _run(g) - assert tuple(o.shape) == (g["b"], g["block"], g["dim"]) - assert torch.isfinite(o.float()).all() - torch.testing.assert_close(o, _run(g)) # deterministic - - -def test_attention_forward_does_not_mutate_kv_cache(): - """The rolling window write must be functional (cache cloned, not mutated).""" - g = _make_attn_inputs() - before = g["kv_cache0"].clone() - _run(g) - torch.testing.assert_close(g["kv_cache0"], before) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py deleted file mode 100644 index f66a37834fab..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py +++ /dev/null @@ -1,252 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""CUDA-graph-safety unit tests for the DSpark batched draft-attention path. - -The load-bearing invariant for the DSpark batched draft path (the default, used -whenever ``cuda_graph_config`` is set, since the one-engine drafter is captured in -the target's graph) is that the batched, sync-free primitives are **numerically -identical, per request**, to the validated scalar path — only the host-int -``start_pos`` and the per-request window indexing are tensorized. These tests assert -that equivalence on CPU (so they run in pre-merge CI without a GPU), plus a -GPU-gated capture+replay smoke test that proves the batched attention is actually -graph-capturable. -""" - -import pytest -import torch - -from tensorrt_llm._torch.models.dspark.attention import ( - apply_dspark_rotary, - apply_dspark_rotary_batched, - dspark_attention_forward, - dspark_attention_forward_batched, - get_dspark_topk_idxs, - get_dspark_topk_idxs_batched, - precompute_dspark_freqs_cis, -) - - -def _make_batched_inputs(seed=0, start_positions=(1, 3, 20)): - """Per-request DSpark attention inputs/weights (shared weights, distinct pos). - - Mirrors ``test_dspark_attention._make_attn_inputs`` but builds ``G`` requests - each with its own ``start_pos`` (small => partial context, large => full - rolling window) and its own pre-seeded window, so the batched-vs-scalar - comparison exercises the per-request RoPE gather + windowed context read. - """ - torch.manual_seed(seed) - dim, n_heads, head_dim, rd = 12, 4, 8, 4 - q_lora, o_lora, n_groups = 6, 5, 2 - window, block = 8, 3 - G = len(start_positions) - bf = torch.bfloat16 - # A single fixed RoPE table covering every request's positions (both paths - # index/gather the same values, so freqs are identical across paths). - maxlen = max(start_positions) + 1 + block + 4 - g = dict( - dim=dim, - n_heads=n_heads, - head_dim=head_dim, - rope_head_dim=rd, - q_lora=q_lora, - o_lora=o_lora, - n_groups=n_groups, - window=window, - block=block, - G=G, - start_positions=list(start_positions), - eps=1e-6, - softmax_scale=head_dim**-0.5, - x=torch.randn(G, block, dim, dtype=bf), - main_x=torch.randn(G, 1, dim, dtype=bf), - # Distinct, non-zero seeded window per request to exercise context reads. - kv_cache=torch.randn(G, window, head_dim, dtype=bf), - wq_a=torch.randn(q_lora, dim, dtype=bf) * 0.1, - wq_b=torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1, - wkv=torch.randn(head_dim, dim, dtype=bf) * 0.1, - wo_a=torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1, - wo_b=torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1, - q_norm=torch.ones(q_lora), - kv_norm=torch.ones(head_dim), - attn_sink=torch.randn(n_heads), - freqs=precompute_dspark_freqs_cis(rd, maxlen), - ) - return g - - -def _attn_kwargs(g): - return dict( - wq_a=g["wq_a"], - q_norm_w=g["q_norm"], - wq_b=g["wq_b"], - wkv=g["wkv"], - kv_norm_w=g["kv_norm"], - wo_a=g["wo_a"], - wo_b=g["wo_b"], - attn_sink=g["attn_sink"], - n_heads=g["n_heads"], - head_dim=g["head_dim"], - rope_head_dim=g["rope_head_dim"], - n_groups=g["n_groups"], - o_lora_rank=g["o_lora"], - window_size=g["window"], - eps=g["eps"], - softmax_scale=g["softmax_scale"], - freqs_cis=g["freqs"], - ) - - -def _scalar_reference(g): - """Run the validated scalar attention once per request and stack the outputs.""" - outs = [] - for i, sp in enumerate(g["start_positions"]): - out_i = dspark_attention_forward( - g["x"][i : i + 1], - g["main_x"][i : i + 1], - int(sp), - g["kv_cache"][i : i + 1].clone(), - **_attn_kwargs(g), - ) - outs.append(out_i) - return torch.cat(outs, dim=0) - - -def _batched(g, persist=False): - G = g["G"] - start_pos = torch.tensor(g["start_positions"], dtype=torch.long) - slots = torch.arange(G, dtype=torch.long) - return dspark_attention_forward_batched( - g["x"], - g["main_x"], - start_pos, - g["kv_cache"].clone(), - slots, - persist=persist, - **_attn_kwargs(g), - ) - - -@pytest.mark.parametrize("seed", [0, 1, 2]) -def test_batched_attention_matches_scalar_per_request(seed): - """Batched attention == per-request scalar attention at distinct start_pos. - - This is the invariant that lets the batched path replace the per-request loop - under CUDA graphs without changing draft quality / greedy parity. - """ - g = _make_batched_inputs(seed=seed) - ref = _scalar_reference(g) - got = _batched(g) - assert tuple(got.shape) == (g["G"], g["block"], g["dim"]) - torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) - - -def test_batched_attention_persist_writes_through_window(): - """persist=True writes main_kv into the shared window at start_pos%window.""" - g = _make_batched_inputs(seed=3) - G, win = g["G"], g["window"] - start_pos = torch.tensor(g["start_positions"], dtype=torch.long) - slots = torch.arange(G, dtype=torch.long) - cache = g["kv_cache"].clone() - before = cache.clone() - dspark_attention_forward_batched( - g["x"], g["main_x"], start_pos, cache, slots, persist=True, **_attn_kwargs(g) - ) - # Exactly the start_pos%win row of each request changed. - for i, sp in enumerate(g["start_positions"]): - changed = (cache[i] != before[i]).any(dim=-1) - expected = torch.zeros(win, dtype=torch.bool) - expected[sp % win] = True - assert torch.equal(changed, expected), f"req {i}: wrong window row written" - - -def test_batched_attention_no_persist_keeps_window(): - """persist=False must not mutate the caller's window (functional).""" - g = _make_batched_inputs(seed=4) - before = g["kv_cache"].clone() - _batched(g, persist=False) - torch.testing.assert_close(g["kv_cache"], before) - - -@pytest.mark.parametrize("start_positions", [(1, 3, 20), (5, 5, 5), (2, 7, 200)]) -def test_batched_topk_matches_scalar(start_positions): - """Fixed-size masked batched topk == scalar topk per request (valid slots).""" - window, block = 8, 3 - start_pos = torch.tensor(start_positions, dtype=torch.long) - batched = get_dspark_topk_idxs_batched(window, block, start_pos) - assert tuple(batched.shape) == (len(start_positions), block, window + block) - for i, sp in enumerate(start_positions): - scalar = get_dspark_topk_idxs(window, 1, block, int(sp))[0] # [block, topk_i] - # The batched row drops the -1-masked context slots to recover the scalar - # (variable-width) index list; both must then be identical. - for m in range(block): - valid = batched[i, m][batched[i, m] >= 0] - torch.testing.assert_close(valid, scalar[m].to(valid.dtype)) - - -@pytest.mark.parametrize("ndim", [3, 4]) -def test_batched_rotary_matches_scalar_per_row(ndim): - """apply_dspark_rotary_batched (per-row freqs) == scalar applied row by row.""" - torch.manual_seed(0) - G, s, h, rd = 3, 4, 2, 8 - x = torch.randn(G, s, h, rd) if ndim == 4 else torch.randn(G, s, rd) - table = precompute_dspark_freqs_cis(rd, 64) - # Per-row absolute start positions -> per-row freq windows. - starts = [1, 9, 30] - per_row = torch.stack([table[sp : sp + s] for sp in starts], dim=0) # [G, s, rd/2] - got = apply_dspark_rotary_batched(x, per_row) - for i, sp in enumerate(starts): - ref_i = apply_dspark_rotary(x[i : i + 1], table[sp : sp + s]) - torch.testing.assert_close(got[i : i + 1], ref_i) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") -def test_batched_attention_cuda_graph_capture_replay(): - """The batched attention captures + replays and matches eager output. - - Proves the path is free of capture-illegal ops (host syncs, dynamic shapes). - """ - g = _make_batched_inputs(seed=0) - dev = "cuda" - G = g["G"] - start_pos = torch.tensor(g["start_positions"], dtype=torch.long, device=dev) - slots = torch.arange(G, dtype=torch.long, device=dev) - # Static input tensors the graph reads/writes. - x = g["x"].to(dev) - main_x = g["main_x"].to(dev) - cache = g["kv_cache"].to(dev) - kw = {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in _attn_kwargs(g).items()} - - def run(persist): - return dspark_attention_forward_batched( - x, main_x, start_pos, cache, slots, persist=persist, **kw - ) - - eager = run(persist=False) - - # Warmup (PyTorch CUDA-graph semantics) then capture on a non-persist call so - # the comparison isn't perturbed by the window write-through. - s = torch.cuda.Stream() - s.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(s): - for _ in range(3): - run(persist=False) - torch.cuda.current_stream().wait_stream(s) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = run(persist=False) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(out, eager, rtol=2e-2, atol=2e-2) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py deleted file mode 100644 index fed0da3a2803..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py +++ /dev/null @@ -1,99 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Unit tests for the DSpark draft I/O proposal stage.""" - -import torch - -from tensorrt_llm._torch.models.dspark.draft import build_draft_input_ids, dspark_propose -from tensorrt_llm._torch.models.dspark.heads import DSparkConfidenceHead, build_markov_head - -VOCAB, HID, RANK, B, BLK = 257, 32, 16, 2, 5 -NOISE_ID = 199 - - -def test_build_draft_input_ids(): - bonus = torch.tensor([7, 9]) - ids = build_draft_input_ids(bonus, block_size=BLK, noise_token_id=NOISE_ID) - assert ids.shape == (B, BLK) - assert torch.equal(ids[:, 0], bonus) - assert torch.all(ids[:, 1:] == NOISE_ID) - - -def test_dspark_propose_full_block_no_confidence(): - torch.manual_seed(0) - markov = build_markov_head( - markov_head_type="rnn", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID - ).eval() - base = torch.randn(B, BLK, VOCAB) - bonus = torch.randint(0, VOCAB, (B,)) - hid = torch.randn(B, BLK, HID) - with torch.no_grad(): - tokens, num = dspark_propose( - base, - bonus_token_ids=bonus, - block_hidden=hid, - markov_head=markov, - confidence_head=None, - block_size=BLK, - ) - assert tokens.shape == (B, BLK) - # No confidence head -> propose the full block. - assert torch.all(num == BLK) - # Tokens match the markov head's own greedy block sampling. - ref_tokens, _ = markov.sample_block_tokens( - base, first_prev_token_ids=bonus, hidden_states=hid, temperature=0.0 - ) - assert torch.equal(tokens, ref_tokens) - - -def test_dspark_propose_confidence_truncates(): - torch.manual_seed(1) - markov = build_markov_head( - markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID - ).eval() - conf = DSparkConfidenceHead(hidden_size=HID).eval() - # The confidence proj is bias-free, so drive the logit via a constant weight - # against a constant hidden: logit = weight_val * HID per position. - base = torch.randn(1, BLK, VOCAB) - bonus = torch.randint(0, VOCAB, (1,)) - hid = torch.ones(1, BLK, HID) - with torch.no_grad(): - conf.proj.weight.fill_(5.0 / HID) # logit ~ 5 -> sigmoid ~ 0.993, all confident - with torch.no_grad(): - _, num = dspark_propose( - base, - bonus_token_ids=bonus, - block_hidden=hid, - markov_head=markov, - confidence_head=conf, - block_size=BLK, - confidence_threshold=0.5, - ) - # All-confident -> full block proposed. - assert int(num[0]) == BLK - # Now make the head output low confidence everywhere -> truncate to 0... but - # confident_prefix_length returns first sub-threshold index (0 here). - with torch.no_grad(): - conf.proj.weight.fill_(-5.0 / HID) # logit ~ -5 -> sigmoid ~ 0.0067 < 0.5 - _, num2 = dspark_propose( - base, - bonus_token_ids=bonus, - block_hidden=hid, - markov_head=markov, - confidence_head=conf, - block_size=BLK, - confidence_threshold=0.5, - ) - assert int(num2[0]) == 0 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py deleted file mode 100644 index 16b6a01acd17..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py +++ /dev/null @@ -1,117 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Unit tests for the DSpark draft-network heads (hardware-agnostic, CPU).""" - -import pytest -import torch - -from tensorrt_llm._torch.models.dspark.heads import ( - DSparkConfidenceHead, - RNNHead, - VanillaMarkov, - build_markov_head, - confident_prefix_length, -) - -VOCAB, RANK, HID, B, BLK = 257, 16, 32, 3, 5 - - -@pytest.mark.parametrize("head_type", ["vanilla", "gated", "rnn"]) -def test_markov_block_sampling_shapes_and_determinism(head_type): - torch.manual_seed(0) - head = build_markov_head( - markov_head_type=head_type, vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID - ).eval() - base = torch.randn(B, BLK, VOCAB) - first = torch.randint(0, VOCAB, (B,)) - hid = torch.randn(B, BLK, HID) - with torch.no_grad(): - tok, logits = head.sample_block_tokens( - base, first_prev_token_ids=first, hidden_states=hid, temperature=0.0 - ) - tok2, _ = head.sample_block_tokens( - base, first_prev_token_ids=first, hidden_states=hid, temperature=0.0 - ) - assert tok.shape == (B, BLK) - assert logits.shape == (B, BLK, VOCAB) - # Greedy is deterministic. - assert torch.equal(tok, tok2) - # Each sampled token is the argmax of its (bias-corrected) step logits. - assert torch.equal(tok, logits.argmax(dim=-1)) - - -def test_markov_bias_is_additive_low_rank(): - # bias = W2(W1[token]); the corrected first-step logits == base + bias. - torch.manual_seed(1) - head = VanillaMarkov(vocab_size=VOCAB, markov_rank=RANK).eval() - base = torch.randn(B, BLK, VOCAB) - first = torch.randint(0, VOCAB, (B,)) - with torch.no_grad(): - _, corrected = head.sample_block_tokens( - base, first_prev_token_ids=first, hidden_states=None, temperature=0.0 - ) - expected0 = base[:, 0] + head.markov_w2(head.markov_w1(first)) - assert torch.allclose(corrected[:, 0], expected0, atol=1e-5) - - -def test_rnn_state_carries_across_positions(): - torch.manual_seed(2) - head = RNNHead(vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID).eval() - initial_state = torch.zeros(1, RANK) - prev_embedding = head.get_prev_embeddings(torch.zeros(1, dtype=torch.long)) - prefix_hidden = torch.randn(1, HID) - current_hidden = torch.randn(1, HID) - - with torch.no_grad(): - state_a, _ = head._rnn_step(initial_state, prev_embedding, prefix_hidden) - state_b, _ = head._rnn_step(initial_state, prev_embedding, -prefix_hidden) - _, bias_a = head._rnn_step(state_a, prev_embedding, current_hidden) - _, bias_b = head._rnn_step(state_b, prev_embedding, current_hidden) - - assert not torch.allclose(state_a, state_b) - assert not torch.allclose(bias_a, bias_b) - - -def test_build_markov_head_rank_zero_returns_none(): - assert ( - build_markov_head( - markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=0, hidden_size=HID - ) - is None - ) - - -def test_confidence_head_and_prefix_truncation(): - head = DSparkConfidenceHead(hidden_size=HID) - conf = head(torch.randn(B, BLK, HID)) - assert conf.shape == (B, BLK) - # threshold 0 disables truncation. - assert confident_prefix_length(conf, block_size=BLK, threshold=0.0) == BLK - # First sub-threshold position truncates the prefix. - logits = torch.tensor([[10.0, 10.0, -10.0, 10.0, 10.0]]) - assert confident_prefix_length(logits, block_size=BLK, threshold=0.5) == 2 - # All-confident -> full block. - logits_hi = torch.full((1, BLK), 10.0) - assert confident_prefix_length(logits_hi, block_size=BLK, threshold=0.5) == BLK - - -def test_confidence_head_with_markov_concat_dim(): - head = DSparkConfidenceHead(hidden_size=HID, markov_rank=RANK, with_markov=True) - hid = torch.randn(B, BLK, HID) - prev_emb = torch.randn(B, BLK, RANK) - out = head(hid, prev_embeddings=prev_emb) - assert out.shape == (B, BLK) - with pytest.raises(AssertionError): - head(hid) # with_markov requires prev_embeddings diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model.py new file mode 100644 index 000000000000..7b87d5375a2e --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model.py @@ -0,0 +1,1036 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the DSpark draft model: captured-context attention primitives, +the batched CUDA-graph-safe attention path, the draft I/O proposal stage, and +the draft-network heads (markov + confidence). CPU except per-test GPU gates. +""" + +import types +from unittest.mock import Mock + +import pytest +import torch +import torch.nn.functional as F + +import tensorrt_llm._torch.models.modeling_dspark as modeling_dspark +from tensorrt_llm._torch.models.dspark.attention import ( + apply_dspark_rotary, + apply_dspark_rotary_batched, + dspark_attention_forward, + dspark_attention_forward_batched, + dspark_sparse_attn, + get_dspark_topk_idxs, + get_dspark_topk_idxs_batched, + precompute_dspark_freqs_cis, +) +from tensorrt_llm._torch.models.dspark.draft import build_draft_input_ids, dspark_propose +from tensorrt_llm._torch.models.dspark.heads import ( + DSparkConfidenceHead, + RNNHead, + VanillaMarkov, + build_markov_head, +) +from tensorrt_llm._torch.models.modeling_dspark import DSparkDraftModel + +VOCAB, RANK, HID, BLK = 257, 16, 32, 5 +DRAFT_B, HEADS_B = 2, 3 +NOISE_ID = 199 + +# --------------------------------------------------------------------------- +# Captured-context attention primitives (CPU) +# --------------------------------------------------------------------------- + + +def test_rope_table_is_cached_once_per_device(): + model = types.SimpleNamespace( + _attn_params={"rope_head_dim": 16}, + _freqs_cap=64, + _rope_theta=10000.0, + _freqs_table_cache={}, + ) + + first = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + second = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + + assert first.data_ptr() == second.data_ptr() + assert len(model._freqs_table_cache) == 1 + positions = torch.tensor([1, 17, 63]) + expected = precompute_dspark_freqs_cis(16, 64, rope_theta=10000.0) + torch.testing.assert_close(first[positions], expected[positions]) + + +def test_dspark_block_uses_stage_id_as_attention_layer_idx(monkeypatch): + captured = {} + + def fake_decoder_layer_init( + self, + model_config, + layer_idx, + aux_stream_dict, + attention_layer_idx=None, + mapping_with_cp=None, + disable_post_moe_fusion=False, + ): + torch.nn.Module.__init__(self) + self.model_config = model_config + self.config = model_config.pretrained_config + self.layer_idx = layer_idx + captured.update( + layer_idx=layer_idx, + attention_layer_idx=attention_layer_idx, + aux_stream_dict=aux_stream_dict, + mapping_with_cp=mapping_with_cp, + disable_post_moe_fusion=disable_post_moe_fusion, + ) + + monkeypatch.setattr( + modeling_dspark.DeepseekV4DecoderLayer, + "__init__", + fake_decoder_layer_init, + ) + model_config = types.SimpleNamespace( + pretrained_config=types.SimpleNamespace(vocab_size=128, hc_mult=2), + spec_config=None, + ) + + block = modeling_dspark.DSparkBlock( + model_config, + layer_idx=10, + aux_stream_dict={}, + stage_id=1, + num_stages=3, + num_capture_layers=0, + ) + + assert block.layer_idx == captured["layer_idx"] == 10 + assert captured["attention_layer_idx"] == block.stage_id == 1 + assert captured["disable_post_moe_fusion"] is True + + +@pytest.mark.parametrize("enable_fused_hc", [True, False]) +def test_forward_stage_honors_enable_fused_hc(monkeypatch, enable_fused_hc): + """The draft stage must use the inherited fused-HC rollback setting.""" + torch.manual_seed(71) + num_requests, block_size, hc_mult, hidden_size = 1, 2, 2, 3 + h = torch.randn(num_requests, block_size, hc_mult, hidden_size) + attention_input = torch.randn(num_requests, block_size, hidden_size) + attention_output = torch.randn_like(attention_input) + mid_residual = torch.randn_like(h) + attention_post_mix = torch.randn(num_requests, block_size, hc_mult, 1) + attention_comb_mix = torch.randn(num_requests, block_size, hc_mult, hc_mult) + ffn_post_mix = torch.randn_like(attention_post_mix) + ffn_comb_mix = torch.randn_like(attention_comb_mix) + raw_ffn_input = torch.randn_like(attention_input) + normed_ffn_input = torch.randn_like(attention_input) + moe_output = torch.randn(num_requests * block_size, hidden_size) + final_h = torch.randn_like(h) + events = [] + + def record(name, result): + def call(*args, **kwargs): + events.append(name) + return result + + return call + + monkeypatch.setattr( + modeling_dspark, + "dspark_attention_forward", + Mock(return_value=attention_output), + ) + + hc_attn = types.SimpleNamespace( + pre_mapping=Mock(return_value=(attention_post_mix, attention_comb_mix, attention_input)), + post_mapping=Mock(side_effect=record("attention_post", mid_residual)), + ) + hc_ffn = types.SimpleNamespace( + fused_hc=Mock( + side_effect=record( + "fused", + (mid_residual, ffn_post_mix, ffn_comb_mix, normed_ffn_input), + ) + ), + pre_mapping=Mock( + side_effect=record("ffn_pre", (ffn_post_mix, ffn_comb_mix, raw_ffn_input)) + ), + post_mapping=Mock(side_effect=record("ffn_post", final_h)), + ) + post_attention_layernorm = Mock(side_effect=record("ffn_norm", normed_ffn_input)) + post_attention_layernorm.weight = torch.ones(hidden_size) + post_attention_layernorm.variance_epsilon = 1e-6 + stage = types.SimpleNamespace( + enable_fused_hc=enable_fused_hc, + hc_attn=hc_attn, + hc_ffn=hc_ffn, + input_layernorm=Mock(side_effect=lambda tensor: tensor), + post_attention_layernorm=post_attention_layernorm, + mlp=Mock(return_value=moe_output), + _dspark_attn={}, + ) + model = types.SimpleNamespace( + use_real_mla=False, + _attn_params={"window_size": 2, "head_dim": 1}, + model_config=types.SimpleNamespace( + mapping=types.SimpleNamespace(enable_attention_dp=False, tp_size=8) + ), + ) + + actual = DSparkDraftModel._forward_stage( + model, + stage, + h, + torch.randn(num_requests, hidden_size), + 1, + torch.empty(0), + torch.zeros(num_requests, block_size, dtype=torch.long), + ) + + assert actual is final_h + torch.testing.assert_close( + stage.mlp.call_args.args[0], + normed_ffn_input.reshape(num_requests * block_size, hidden_size), + ) + # Non-attention-DP multi-GPU (enable_attention_dp=False, tp_size>1): the draft + # MoE must all-reduce its TP-sharded output, mirroring the target MoE. The + # attention-DP and single-GPU paths keep it disabled. + assert stage.mlp.call_args.kwargs["final_all_reduce_params"].enable_allreduce is True + if enable_fused_hc: + assert events == ["fused", "ffn_post"] + hc_ffn.fused_hc.assert_called_once() + hc_attn.post_mapping.assert_not_called() + hc_ffn.pre_mapping.assert_not_called() + post_attention_layernorm.assert_not_called() + fused_kwargs = hc_ffn.fused_hc.call_args.kwargs + assert fused_kwargs["norm_weight"] is post_attention_layernorm.weight + assert fused_kwargs["norm_eps"] == post_attention_layernorm.variance_epsilon + else: + assert events == ["attention_post", "ffn_pre", "ffn_norm", "ffn_post"] + hc_ffn.fused_hc.assert_not_called() + hc_attn.post_mapping.assert_called_once() + hc_ffn.pre_mapping.assert_called_once_with(mid_residual) + post_attention_layernorm.assert_called_once_with(raw_ffn_input) + + +def _ref_precompute_freqs_cis(dim, seqlen, base): + """DeepSpec precompute_freqs_cis with original_seq_len == 0 (no YaRN).""" + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + t = torch.arange(seqlen) + freqs = torch.outer(t, freqs) + return torch.polar(torch.ones_like(freqs), freqs) + + +def _ref_apply_rotary_emb(x, freqs_cis, inverse=False): + """DeepSpec apply_rotary_emb (returns a fresh tensor instead of in-place).""" + xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if xc.ndim == 3: + fc = freqs_cis.view(1, xc.size(1), xc.size(-1)) + else: + fc = freqs_cis.view(1, xc.size(1), 1, xc.size(-1)) + return torch.view_as_real(xc * fc).flatten(-2).to(x.dtype) + + +def _loop_reference(q, kv, attn_sink, topk_idxs, scale): + """Obvious, slow per-(b,m,h) reference of the exact kernel math.""" + b, m, h, d = q.shape + out = torch.zeros(b, m, h, d, dtype=torch.float32) + qf, kvf, sink = q.float(), kv.float(), attn_sink.float() + for bi in range(b): + for mi in range(m): + idxs = topk_idxs[bi, mi].tolist() + for hi in range(h): + scores, vecs = [], [] + for j in idxs: + if j < 0: + continue + k = kvf[bi, j] + scores.append(torch.dot(qf[bi, mi, hi], k) * scale) + vecs.append(k) + if not scores: + continue + s = torch.stack(scores) + smax = s.max() + p = torch.exp(s - smax) + denom = p.sum() + torch.exp(sink[hi] - smax) + num = (p.unsqueeze(-1) * torch.stack(vecs)).sum(0) + out[bi, mi, hi] = num / denom + return out + + +def test_sparse_attn_matches_loop_reference(): + torch.manual_seed(0) + b, m, h, d, n, topk = 2, 5, 3, 16, 40, 12 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, n, d) + attn_sink = torch.randn(h) + idx = torch.stack( + [torch.stack([torch.randperm(n)[:topk] for _ in range(m)]) for _ in range(b)] + ).int() + scale = d**-0.5 + got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() + ref = _loop_reference(q, kv, attn_sink, idx, scale) + torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) + + +def test_sparse_attn_no_sink_matches_sdpa(): + """With a -inf sink and all-valid contiguous indices, the primitive must + equal standard scaled-dot-product attention over the gathered KV.""" + torch.manual_seed(0) + b, m, h, d, topk = 2, 4, 5, 16, 9 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, topk, d) # n == topk, attend to all + attn_sink = torch.full((h,), float("-inf")) + idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() + scale = d**-0.5 + got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() + + # SDPA: q [b,h,m,d], k/v [b,h,topk,d] (broadcast the shared KV over heads). + qh = q.permute(0, 2, 1, 3) + kvh = kv.unsqueeze(1).expand(b, h, topk, d) + ref = F.scaled_dot_product_attention(qh, kvh, kvh, scale=scale).permute(0, 2, 1, 3) + torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) + + +def test_sparse_attn_masked_indices_excluded(): + """An index of -1 must be excluded exactly (equiv. to dropping that column).""" + torch.manual_seed(0) + b, m, h, d = 1, 1, 2, 16 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, 5, d) + sink = torch.full((h,), float("-inf")) + scale = d**-0.5 + full = torch.tensor([[[0, 1, 2, 3]]]).int() + masked = torch.tensor([[[0, 1, 2, -1]]]).int() + drop3 = torch.tensor([[[0, 1, 2]]]).int() + got_masked = dspark_sparse_attn(q, kv, sink, masked, scale) + got_drop = dspark_sparse_attn(q, kv, sink, drop3, scale) + torch.testing.assert_close(got_masked, got_drop, rtol=1e-5, atol=1e-5) + # And masking genuinely changes the result vs attending to position 3. + got_full = dspark_sparse_attn(q, kv, sink, full, scale) + assert not torch.allclose(got_full, got_masked, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("start_pos,window,block", [(3, 128, 5), (10, 4, 5)]) +def test_get_dspark_topk_idxs_matches_reference(start_pos, window, block): + bsz = 3 + got = get_dspark_topk_idxs(window, bsz, block, start_pos) + # Reference formula (DeepSpec get_dspark_topk_idxs). + ctx = torch.arange(min(window, start_pos + 1)) + blk = window + torch.arange(block) + ref_row = torch.cat([ctx, blk]).int() + assert got.shape == (bsz, block, ref_row.numel()) + for bi in range(bsz): + for mi in range(block): + torch.testing.assert_close(got[bi, mi], ref_row) + + +@pytest.mark.parametrize("rope_head_dim,seqlen", [(64, 16), (64, 1), (128, 8)]) +def test_precompute_freqs_cis_matches_reference(rope_head_dim, seqlen): + got = precompute_dspark_freqs_cis(rope_head_dim, seqlen, rope_theta=10000.0) + ref = _ref_precompute_freqs_cis(rope_head_dim, seqlen, 10000.0) + torch.testing.assert_close(got, ref) + + +@pytest.mark.parametrize("ndim", [3, 4]) +def test_apply_rotary_matches_reference(ndim): + torch.manual_seed(0) + b, s, h, rd = 2, 5, 4, 64 + x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) + fc = precompute_dspark_freqs_cis(rd, s) + got = apply_dspark_rotary(x, fc) + ref = _ref_apply_rotary_emb(x, fc) + torch.testing.assert_close(got, ref) + # De-rotation (inverse) must undo the forward rotation (property test). + roundtrip = apply_dspark_rotary(got, fc, inverse=True) + torch.testing.assert_close(roundtrip, x, rtol=1e-5, atol=1e-5) + + +def _make_attn_inputs(seed=0): + """Small synthetic DSpark attention inputs/weights (CPU bf16).""" + torch.manual_seed(seed) + dim, n_heads, head_dim, rd = 12, 4, 8, 4 + q_lora, o_lora, n_groups = 6, 5, 2 + window, block, start_pos = 8, 3, 5 + b = 2 + g = dict( + dim=dim, + n_heads=n_heads, + head_dim=head_dim, + rope_head_dim=rd, + q_lora=q_lora, + o_lora=o_lora, + n_groups=n_groups, + window=window, + block=block, + start_pos=start_pos, + b=b, + eps=1e-6, + softmax_scale=head_dim**-0.5, + ) + bf = torch.bfloat16 + g["x"] = torch.randn(b, block, dim, dtype=bf) + g["main_x"] = torch.randn(b, 1, dim, dtype=bf) + g["kv_cache0"] = torch.randn(b, window, head_dim, dtype=bf) + g["wq_a"] = torch.randn(q_lora, dim, dtype=bf) * 0.1 + g["wq_b"] = torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1 + g["wkv"] = torch.randn(head_dim, dim, dtype=bf) * 0.1 + g["wo_a"] = torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1 + g["wo_b"] = torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1 + g["q_norm"] = torch.ones(q_lora) + g["kv_norm"] = torch.ones(head_dim) + g["attn_sink"] = torch.randn(n_heads) + g["freqs"] = precompute_dspark_freqs_cis(rd, start_pos + 1 + block + 2) + return g + + +def _run(g): + return dspark_attention_forward( + g["x"], + g["main_x"], + g["start_pos"], + g["kv_cache0"], + wq_a=g["wq_a"], + q_norm_w=g["q_norm"], + wq_b=g["wq_b"], + wkv=g["wkv"], + kv_norm_w=g["kv_norm"], + wo_a=g["wo_a"], + wo_b=g["wo_b"], + attn_sink=g["attn_sink"], + n_heads=g["n_heads"], + head_dim=g["head_dim"], + rope_head_dim=g["rope_head_dim"], + n_groups=g["n_groups"], + o_lora_rank=g["o_lora"], + window_size=g["window"], + eps=g["eps"], + softmax_scale=g["softmax_scale"], + freqs_cis=g["freqs"], + ) + + +def test_attention_forward_shape_and_determinism(): + g = _make_attn_inputs() + before = g["kv_cache0"].clone() + o = _run(g) + assert tuple(o.shape) == (g["b"], g["block"], g["dim"]) + assert torch.isfinite(o.float()).all() + torch.testing.assert_close(o, _run(g)) # deterministic + # The rolling window write must be functional (cache cloned, not mutated). + torch.testing.assert_close(g["kv_cache0"], before) + + +# --------------------------------------------------------------------------- +# Batched draft-attention path (CUDA-graph safety) +# --------------------------------------------------------------------------- + + +def _make_batched_inputs(seed=0, start_positions=(1, 3, 20)): + """Per-request DSpark attention inputs/weights (shared weights, distinct + start_pos and pre-seeded windows) for batched-vs-scalar comparison.""" + torch.manual_seed(seed) + dim, n_heads, head_dim, rd = 12, 4, 8, 4 + q_lora, o_lora, n_groups = 6, 5, 2 + window, block = 8, 3 + G = len(start_positions) + bf = torch.bfloat16 + # A single fixed RoPE table covering every request's positions (both paths + # index/gather the same values, so freqs are identical across paths). + maxlen = max(start_positions) + 1 + block + 4 + g = dict( + dim=dim, + n_heads=n_heads, + head_dim=head_dim, + rope_head_dim=rd, + q_lora=q_lora, + o_lora=o_lora, + n_groups=n_groups, + window=window, + block=block, + G=G, + start_positions=list(start_positions), + eps=1e-6, + softmax_scale=head_dim**-0.5, + x=torch.randn(G, block, dim, dtype=bf), + main_x=torch.randn(G, 1, dim, dtype=bf), + # Distinct, non-zero seeded window per request to exercise context reads. + kv_cache=torch.randn(G, window, head_dim, dtype=bf), + wq_a=torch.randn(q_lora, dim, dtype=bf) * 0.1, + wq_b=torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1, + wkv=torch.randn(head_dim, dim, dtype=bf) * 0.1, + wo_a=torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1, + wo_b=torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1, + q_norm=torch.ones(q_lora), + kv_norm=torch.ones(head_dim), + attn_sink=torch.randn(n_heads), + freqs=precompute_dspark_freqs_cis(rd, maxlen), + ) + return g + + +def _attn_kwargs(g): + return dict( + wq_a=g["wq_a"], + q_norm_w=g["q_norm"], + wq_b=g["wq_b"], + wkv=g["wkv"], + kv_norm_w=g["kv_norm"], + wo_a=g["wo_a"], + wo_b=g["wo_b"], + attn_sink=g["attn_sink"], + n_heads=g["n_heads"], + head_dim=g["head_dim"], + rope_head_dim=g["rope_head_dim"], + n_groups=g["n_groups"], + o_lora_rank=g["o_lora"], + window_size=g["window"], + eps=g["eps"], + softmax_scale=g["softmax_scale"], + freqs_cis=g["freqs"], + ) + + +def _scalar_reference(g): + """Run the validated scalar attention once per request and stack the outputs.""" + outs = [] + for i, sp in enumerate(g["start_positions"]): + out_i = dspark_attention_forward( + g["x"][i : i + 1], + g["main_x"][i : i + 1], + int(sp), + g["kv_cache"][i : i + 1].clone(), + **_attn_kwargs(g), + ) + outs.append(out_i) + return torch.cat(outs, dim=0) + + +def _batched(g, persist=False): + G = g["G"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long) + slots = torch.arange(G, dtype=torch.long) + return dspark_attention_forward_batched( + g["x"], + g["main_x"], + start_pos, + g["kv_cache"].clone(), + slots, + persist=persist, + **_attn_kwargs(g), + ) + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_batched_attention_matches_scalar_per_request(seed): + """Batched attention == per-request scalar attention at distinct start_pos.""" + g = _make_batched_inputs(seed=seed) + ref = _scalar_reference(g) + got = _batched(g) + assert tuple(got.shape) == (g["G"], g["block"], g["dim"]) + torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) + + +def test_batched_attention_persist_writes_through_window(): + """persist gates the window write: False mutates nothing, True writes + main_kv into exactly the start_pos%window row of each request.""" + g = _make_batched_inputs(seed=3) + G, win = g["G"], g["window"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long) + slots = torch.arange(G, dtype=torch.long) + cache = g["kv_cache"].clone() + before = cache.clone() + + dspark_attention_forward_batched( + g["x"], g["main_x"], start_pos, cache, slots, persist=False, **_attn_kwargs(g) + ) + # persist=False must not mutate the caller's window (functional). + torch.testing.assert_close(cache, before) + + dspark_attention_forward_batched( + g["x"], g["main_x"], start_pos, cache, slots, persist=True, **_attn_kwargs(g) + ) + # Exactly the start_pos%win row of each request changed. + for i, sp in enumerate(g["start_positions"]): + changed = (cache[i] != before[i]).any(dim=-1) + expected = torch.zeros(win, dtype=torch.bool) + expected[sp % win] = True + assert torch.equal(changed, expected), f"req {i}: wrong window row written" + + +@pytest.mark.parametrize("start_positions", [(1, 3, 20), (5, 5, 5), (2, 7, 200)]) +def test_batched_topk_matches_scalar(start_positions): + """Fixed-size masked batched topk == scalar topk per request (valid slots).""" + window, block = 8, 3 + start_pos = torch.tensor(start_positions, dtype=torch.long) + batched = get_dspark_topk_idxs_batched(window, block, start_pos) + assert tuple(batched.shape) == (len(start_positions), block, window + block) + for i, sp in enumerate(start_positions): + scalar = get_dspark_topk_idxs(window, 1, block, int(sp))[0] # [block, topk_i] + # The batched row drops the -1-masked context slots to recover the scalar + # (variable-width) index list; both must then be identical. + for m in range(block): + valid = batched[i, m][batched[i, m] >= 0] + torch.testing.assert_close(valid, scalar[m].to(valid.dtype)) + + +@pytest.mark.parametrize("ndim", [3, 4]) +def test_batched_rotary_matches_scalar_per_row(ndim): + """apply_dspark_rotary_batched (per-row freqs) == scalar applied row by row.""" + torch.manual_seed(0) + G, s, h, rd = 3, 4, 2, 8 + x = torch.randn(G, s, h, rd) if ndim == 4 else torch.randn(G, s, rd) + table = precompute_dspark_freqs_cis(rd, 64) + # Per-row absolute start positions -> per-row freq windows. + starts = [1, 9, 30] + per_row = torch.stack([table[sp : sp + s] for sp in starts], dim=0) # [G, s, rd/2] + got = apply_dspark_rotary_batched(x, per_row) + for i, sp in enumerate(starts): + ref_i = apply_dspark_rotary(x[i : i + 1], table[sp : sp + s]) + torch.testing.assert_close(got[i : i + 1], ref_i) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +def test_batched_attention_cuda_graph_capture_replay(): + """The batched attention captures + replays and matches eager output, + proving the path is free of capture-illegal ops.""" + g = _make_batched_inputs(seed=0) + dev = "cuda" + G = g["G"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long, device=dev) + slots = torch.arange(G, dtype=torch.long, device=dev) + # Static input tensors the graph reads/writes. + x = g["x"].to(dev) + main_x = g["main_x"].to(dev) + cache = g["kv_cache"].to(dev) + kw = {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in _attn_kwargs(g).items()} + + def run(persist): + return dspark_attention_forward_batched( + x, main_x, start_pos, cache, slots, persist=persist, **kw + ) + + eager = run(persist=False) + + # Warmup (PyTorch CUDA-graph semantics) then capture on a non-persist call so + # the comparison isn't perturbed by the window write-through. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run(persist=False) + torch.cuda.current_stream().wait_stream(s) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = run(persist=False) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(out, eager, rtol=2e-2, atol=2e-2) + + +# Confidence-scheduled verification: the scoring path runs inside the target's +# captured graph, so it must itself be graph-capturable. + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +@pytest.mark.parametrize("with_markov", [False, True]) +def test_confidence_head_is_cuda_graph_capturable(with_markov): + B, BLK, HID, RANK = 4, 5, 32, 16 + head = ( + DSparkConfidenceHead( + hidden_size=HID, markov_rank=RANK, with_markov=with_markov, block_size=BLK + ) + .cuda() + .eval() + ) + hid = torch.randn(B, BLK, HID, device="cuda") + prev = torch.randn(B, BLK, RANK, device="cuda") if with_markov else None + out = torch.empty(B, BLK, device="cuda", dtype=torch.float32) + + def run(): + return head.apply_sts(head(hid, prev_embeddings=prev)) + + # Warm up on a side stream, as CUDA graph capture requires. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out.copy_(run()) + + eager = run() + g.replay() + torch.cuda.synchronize() + assert torch.allclose(out, eager, atol=1e-5), "replay diverged from eager" + assert torch.all((out >= 0.0) & (out <= 1.0)), "calibrated confidence must be a probability" + + # New inputs must flow through on replay -- a graph that captured constants + # would keep returning the old scores and silently freeze the scheduler. + hid.copy_(torch.randn_like(hid)) + fresh = run() + g.replay() + torch.cuda.synchronize() + assert torch.allclose(out, fresh, atol=1e-5), "replay ignored updated inputs" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +def test_sts_table_update_is_visible_to_a_captured_graph(): + """STS must be updated in place; rebinding the buffer changes its data_ptr + and a graph captured earlier would keep reading the old storage.""" + B, BLK, HID = 2, 5, 32 + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK).cuda().eval() + hid = torch.randn(B, BLK, HID, device="cuda") + out = torch.empty(B, BLK, device="cuda", dtype=torch.float32) + + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + head.apply_sts(head(hid)) + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out.copy_(head.apply_sts(head(hid))) + + g.replay() + torch.cuda.synchronize() + before = out.clone() + + storage_before = head.sts_temperatures.data_ptr() + head.load_sts_temperatures(torch.full((BLK,), 4.0)) + assert head.sts_temperatures.data_ptr() == storage_before, ( + "load_sts_temperatures must write in place; rebinding hides it from the graph" + ) + + g.replay() + torch.cuda.synchronize() + assert not torch.allclose(out, before), ( + "the captured graph did not see the new STS temperatures" + ) + assert torch.allclose(out, torch.sigmoid(head(hid) / 4.0), atol=1e-5) + + +def test_apply_sts_accepts_host_logits_from_a_device_head(): + """apply_sts must handle host logits while sts_temperatures lives on the + device: the planner calibrates on CPU-staged rows from a CUDA head.""" + BLK = 5 + head = DSparkConfidenceHead(hidden_size=8, block_size=BLK) + if torch.cuda.is_available(): + head = head.cuda() + head.load_sts_temperatures(torch.full((BLK,), 2.0)) + + host_logits = torch.zeros(3, BLK) # CPU, as the planner stages them + probs = head.apply_sts(host_logits) + assert probs.device.type == "cpu" + assert torch.allclose(probs, torch.full((3, BLK), 0.5), atol=1e-6) + + +# Ragged verification runs inside the same captured graph, so its scatter and +# accept-count primitives must also capture without host syncs. + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +def test_ragged_scatter_is_cuda_graph_capturable(): + """scatter_ragged_to_padded and count_accepted_ragged capture and replay + inside the target's graph without host syncs.""" + from tensorrt_llm._torch.speculative.dspark_ragged import ( + build_qo_indptr, + count_accepted_ragged, + scatter_ragged_to_padded, + ) + + lens = torch.tensor([3, 1, 2], dtype=torch.int32, device="cuda") + indptr = build_qo_indptr(lens) + flat = torch.tensor([10, 11, 12, 20, 30, 31], dtype=torch.int32, device="cuda") + draft = torch.tensor([[1, 2, 3, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") + target = torch.tensor([[1, 2, 7, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") + out = torch.empty(3, 4, dtype=torch.int32, device="cuda") + accepted = torch.empty(3, dtype=torch.int64, device="cuda") + + def run(): + padded = scatter_ragged_to_padded( + flat, verify_lens=lens, qo_indptr=indptr, max_len=4, pad_value=-1 + ) + counts = count_accepted_ragged( + draft_tokens=draft, target_tokens=target, verify_lens=lens + ) + return padded, counts + + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + padded, counts = run() + out.copy_(padded) + accepted.copy_(counts) + + g.replay() + torch.cuda.synchronize() + assert out.tolist() == [[10, 11, 12, -1], [20, -1, -1, -1], [30, 31, -1, -1]] + # req 0 matches 2 then diverges; req 1 matches its single position; req 2 + # matches both. Column 3 is padding for every row and must never count. + assert accepted.tolist() == [2, 1, 2] + + # New packed contents must flow through on replay. + flat.copy_(torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32, device="cuda")) + g.replay() + torch.cuda.synchronize() + assert out.tolist() == [[1, 2, 3, -1], [4, -1, -1, -1], [5, 6, -1, -1]] + + +# --------------------------------------------------------------------------- +# Draft I/O proposal stage +# --------------------------------------------------------------------------- + + +def test_build_draft_input_ids(): + bonus = torch.tensor([7, 9]) + ids = build_draft_input_ids(bonus, block_size=BLK, noise_token_id=NOISE_ID) + assert ids.shape == (DRAFT_B, BLK) + assert torch.equal(ids[:, 0], bonus) + assert torch.all(ids[:, 1:] == NOISE_ID) + + +def test_dspark_propose_full_block_no_confidence(): + torch.manual_seed(0) + markov = build_markov_head( + markov_head_type="rnn", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID + ).eval() + base = torch.randn(DRAFT_B, BLK, VOCAB) + bonus = torch.randint(0, VOCAB, (DRAFT_B,)) + hid = torch.randn(DRAFT_B, BLK, HID) + with torch.no_grad(): + tokens, confidence = dspark_propose( + base, + bonus_token_ids=bonus, + block_hidden=hid, + markov_head=markov, + confidence_head=None, + block_size=BLK, + ) + assert tokens.shape == (DRAFT_B, BLK) + assert confidence is None + # Tokens match the markov head's own greedy block sampling. + ref_tokens, _ = markov.sample_block_tokens( + base, first_prev_token_ids=bonus, hidden_states=hid, temperature=0.0 + ) + assert torch.equal(tokens, ref_tokens) + + +def _propose_with_confidence(conf, markov, *, return_confidence): + base = torch.randn(1, BLK, VOCAB) + bonus = torch.randint(0, VOCAB, (1,)) + hid = torch.ones(1, BLK, HID) + with torch.no_grad(): + return dspark_propose( + base, + bonus_token_ids=bonus, + block_hidden=hid, + markov_head=markov, + confidence_head=conf, + block_size=BLK, + return_confidence=return_confidence, + ) + + +def test_dspark_propose_scores_without_shortening_the_block(): + """The block is always proposed in full; confidence only scores it, and + only when asked (return_confidence is the run-constant opt-in).""" + torch.manual_seed(1) + markov = build_markov_head( + markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID + ).eval() + conf = DSparkConfidenceHead(hidden_size=HID, block_size=BLK).eval() + # The confidence proj is bias-free, so drive the logit via a constant weight + # against a constant hidden: logit = weight_val * HID per position. + with torch.no_grad(): + conf.proj.weight.fill_(-5.0 / HID) # sigmoid ~ 0.0067, i.e. hopeless + + tokens, confidence = _propose_with_confidence(conf, markov, return_confidence=True) + assert tokens.shape == (1, BLK) + assert confidence.shape == (1, BLK) + # Low confidence must NOT shorten the proposal -- that decision belongs to + # the verification scheduler, not the drafter. + assert torch.all(confidence < 0.0) + assert torch.all(conf.apply_sts(confidence) < 0.5) + + # Opt-in: with the flag off the head is not consulted at all. + _, confidence = _propose_with_confidence(conf, markov, return_confidence=False) + assert confidence is None + + +def test_dspark_propose_is_free_of_host_syncs(): + """No ``.item()``/``nonzero`` on this path: it runs inside the target's graph.""" + import inspect + import io + import tokenize + + from tensorrt_llm._torch.models.dspark import draft as draft_mod + + src = inspect.getsource(draft_mod.dspark_propose) + # Strip comments and string literals so the check reads the code, not the + # prose describing it (the implementation comments name these very calls). + code = "".join( + tok.string if tok.type not in (tokenize.COMMENT, tokenize.STRING) else " " + for tok in tokenize.generate_tokens(io.StringIO(src).readline) + ) + for banned in (".item(", "nonzero", "range("): + assert banned not in code, f"dspark_propose must stay capture-safe: found {banned!r}" + + +# --------------------------------------------------------------------------- +# Draft-network heads (markov + confidence, CPU) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("head_type", ["vanilla", "gated", "rnn"]) +def test_markov_block_sampling_shapes_and_determinism(head_type): + torch.manual_seed(0) + head = build_markov_head( + markov_head_type=head_type, vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID + ).eval() + base = torch.randn(HEADS_B, BLK, VOCAB) + first = torch.randint(0, VOCAB, (HEADS_B,)) + hid = torch.randn(HEADS_B, BLK, HID) + with torch.no_grad(): + tok, logits = head.sample_block_tokens( + base, first_prev_token_ids=first, hidden_states=hid, temperature=0.0 + ) + tok2, _ = head.sample_block_tokens( + base, first_prev_token_ids=first, hidden_states=hid, temperature=0.0 + ) + assert tok.shape == (HEADS_B, BLK) + assert logits.shape == (HEADS_B, BLK, VOCAB) + # Greedy is deterministic. + assert torch.equal(tok, tok2) + # Each sampled token is the argmax of its (bias-corrected) step logits. + assert torch.equal(tok, logits.argmax(dim=-1)) + + +def test_markov_bias_is_additive_low_rank(): + # bias = W2(W1[token]); the corrected first-step logits == base + bias. + torch.manual_seed(1) + head = VanillaMarkov(vocab_size=VOCAB, markov_rank=RANK).eval() + base = torch.randn(HEADS_B, BLK, VOCAB) + first = torch.randint(0, VOCAB, (HEADS_B,)) + with torch.no_grad(): + _, corrected = head.sample_block_tokens( + base, first_prev_token_ids=first, hidden_states=None, temperature=0.0 + ) + expected0 = base[:, 0] + head.markov_w2(head.markov_w1(first)) + assert torch.allclose(corrected[:, 0], expected0, atol=1e-5) + + +def test_rnn_state_carries_across_positions(): + torch.manual_seed(2) + head = RNNHead(vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID).eval() + initial_state = torch.zeros(1, RANK) + prev_embedding = head.get_prev_embeddings(torch.zeros(1, dtype=torch.long)) + prefix_hidden = torch.randn(1, HID) + current_hidden = torch.randn(1, HID) + + with torch.no_grad(): + state_a, _ = head._rnn_step(initial_state, prev_embedding, prefix_hidden) + state_b, _ = head._rnn_step(initial_state, prev_embedding, -prefix_hidden) + _, bias_a = head._rnn_step(state_a, prev_embedding, current_hidden) + _, bias_b = head._rnn_step(state_b, prev_embedding, current_hidden) + + assert not torch.allclose(state_a, state_b) + assert not torch.allclose(bias_a, bias_b) + + +def test_build_markov_head_rank_zero_returns_none(): + assert ( + build_markov_head( + markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=0, hidden_size=HID + ) + is None + ) + + +def test_confidence_head_emits_raw_logits(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + conf = head(torch.randn(HEADS_B, BLK, HID)) + assert conf.shape == (HEADS_B, BLK) + assert conf.dtype == torch.float32 + + +def test_apply_sts_uses_per_position_temperatures(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + # Default (all-ones) calibration is a plain sigmoid, bounded in [0, 1]. + raw = torch.tensor([[2.0, 0.0, -2.0, 1.0, -1.0]]) + assert torch.allclose(head.apply_sts(raw), torch.sigmoid(raw), atol=1e-6) + assert torch.all((head.apply_sts(raw) >= 0.0) & (head.apply_sts(raw) <= 1.0)) + + temps = torch.tensor([1.0, 2.0, 4.0, 1.0, 1.0]) + head.load_sts_temperatures(temps) + raw = torch.full((1, BLK), 4.0) + got = head.apply_sts(raw) + # A larger temperature pulls the probability toward 0.5. + assert got[0, 0] > got[0, 1] > got[0, 2] + assert torch.allclose(got, torch.sigmoid(raw / temps), atol=1e-6) + + +def test_load_sts_temperatures_updates_in_place(): + """The STS table is calibration, not a checkpoint weight: it must update + in place, stay out of state_dict, and refuse malformed tables.""" + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + assert "sts_temperatures" not in head.state_dict() + with pytest.raises(ValueError, match="one per block position"): + head.load_sts_temperatures(torch.ones(BLK + 1)) + with pytest.raises(ValueError, match="strictly positive"): + head.load_sts_temperatures(torch.zeros(BLK)) + before = head.sts_temperatures + head.load_sts_temperatures(torch.full((BLK,), 2.0)) + assert head.sts_temperatures is before + assert torch.allclose(before, torch.full((BLK,), 2.0)) + + +def test_confidence_head_load_weights_rejects_a_dropped_bias(): + """A checkpoint bias silently ignored would shift every confidence score.""" + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + good = {"proj.weight": torch.randn(1, HID)} + head.load_weights([good]) + assert torch.allclose(head.proj.weight, good["proj.weight"]) + + with pytest.raises(ValueError, match="bias=True"): + head.load_weights([{**good, "proj.bias": torch.zeros(1)}]) + + # Constructed with bias=True, the same checkpoint bias is accepted. + biased = DSparkConfidenceHead(hidden_size=HID, block_size=BLK, bias=True) + w = {"proj.weight": torch.randn(1, HID), "proj.bias": torch.randn(1)} + biased.load_weights([w]) + assert torch.allclose(biased.proj.bias, w["proj.bias"]) + + +def test_confidence_head_with_markov_concat_dim(): + head = DSparkConfidenceHead(hidden_size=HID, markov_rank=RANK, with_markov=True) + hid = torch.randn(HEADS_B, BLK, HID) + prev_emb = torch.randn(HEADS_B, BLK, RANK) + out = head(hid, prev_embeddings=prev_emb) + assert out.shape == (HEADS_B, BLK) + with pytest.raises(AssertionError): + head(hid) # with_markov requires prev_embeddings diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_planner.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_planner.py new file mode 100644 index 000000000000..f59f7bad6fb8 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_planner.py @@ -0,0 +1,903 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark planner/budget/cost-table unit tests (hardware-agnostic, CPU). + +Covers survival, the budget argmax, cost-table interpolation, top-k +allocation, tier derivation, SGLang decision parity, the verify-length pin, +snapshot instruments, and reconciliation guards. +""" + +import itertools + +import numpy as np +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, budget_argmax_over_uniform_lens, check_table_fingerprint, + compute_verify_token_budget, derive_verify_len_tiers) +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, compute_survival, schedule_verify_lens_topk) +from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + +BLOCK = 7 + + +def _cfg(**kw) -> DSparkScheduleConfig: + return DSparkScheduleConfig(block_size=BLOCK, **kw) + + +# -------------------------------------------------------------------------- +# survival +# -------------------------------------------------------------------------- + + +def test_survival_is_cumulative_product(): + conf = torch.tensor([[0.9, 0.8, 0.5, 1.0, 1.0, 1.0, 1.0]]) + surv = compute_survival(conf) + assert torch.allclose(surv[0, :3], torch.tensor([0.9, 0.72, 0.36]), atol=1e-6) + + +# -------------------------------------------------------------------------- +# budget argmax +# -------------------------------------------------------------------------- + + +def _brute_force_budget(survival, num_gen, table, min_len=1, max_len=None): + """Independent O(N) reimplementation used as the oracle for the argmax: + expected emitted tokens per millisecond over submitted-token cost.""" + bs, blk = survival.shape + cap = min(max_len or blk, blk) + cand = np.sort(survival[:, min_len:cap].reshape(-1))[::-1] + base = num_gen + float(survival[:, :min_len].sum()) + floor_tokens = bs * (min_len + 1) + best_n, best = 0, -np.inf + for n in range(cand.size + 1): + theta = (base + float(cand[:n].sum())) / table.step_time(floor_tokens + n, num_gen) + if theta > best: + best, best_n = theta, n + return best_n + + +def test_budget_matches_brute_force_on_staircase(): + """Theta is not unimodal on a staircase; a greedy first-descent scan is wrong.""" + rng = np.random.default_rng(7) + # A cost curve with genuine shelves and risers. + table = SpsCostTable( + token_counts=(0, 64, 128, 192, 256, 320), + step_time_ms=(4.0, 4.05, 6.5, 6.55, 9.0, 9.05), + ) + for trial in range(25): + bs = int(rng.integers(1, 40)) + conf = rng.uniform(0.55, 0.99, size=(bs, BLOCK)) + surv = np.cumprod(conf, axis=1) + got = compute_verify_token_budget(survival=surv, num_gen_requests=bs, cost_table=table) + want = _brute_force_budget(surv, bs, table) + assert got == want, f"trial {trial}: bs={bs} got={got} want={want}" + + +def test_flat_cost_table_degenerates_to_verify_all(): + """A flat cost model makes every token free -> spend the whole budget.""" + table = SpsCostTable.flat(2.0) + assert table.is_flat + surv = np.cumprod(np.full((6, BLOCK), 0.9), axis=1) + budget = compute_verify_token_budget(survival=surv, num_gen_requests=6, cost_table=table) + assert budget == 6 * (BLOCK - 1) + # An empty batch has no candidates to buy, whatever the table says. + assert ( + compute_verify_token_budget( + survival=np.zeros((0, BLOCK)), num_gen_requests=0, cost_table=table + ) + == 0 + ) + + +def test_cost_table_interpolates_between_breakpoints(): + """The consumer contract is clamped linear interpolation, not a floor; + a shelf is only flat if the table measured it as two equal breakpoints.""" + table = SpsCostTable(token_counts=(0, 100), step_time_ms=(1.0, 5.0)) + assert table.step_time(99) > table.step_time(1) + assert table.step_time(50) == pytest.approx(3.0) + # A measured shelf stays flat under interpolation. + shelf = SpsCostTable(token_counts=(0, 99, 100), step_time_ms=(1.0, 1.0, 5.0)) + assert shelf.step_time(1) == shelf.step_time(99) + # Above the last breakpoint we clamp rather than extrapolate. + assert table.step_time(10_000) == table.step_time(100) + + +def test_cost_table_rejects_malformed_input(): + with pytest.raises(ValueError, match="strictly increasing"): + SpsCostTable(token_counts=(0, 0), step_time_ms=(1.0, 2.0)) + with pytest.raises(ValueError, match="same length"): + SpsCostTable(token_counts=(0, 1), step_time_ms=(1.0,)) + with pytest.raises(ValueError, match="positive"): + SpsCostTable(token_counts=(0,), step_time_ms=(0.0,)) + + +# -------------------------------------------------------------------------- +# per-request allocation +# -------------------------------------------------------------------------- + + +def test_allocation_is_a_prefix_and_respects_bounds(): + torch.manual_seed(3) + cfg = _cfg(min_verify_len=1) + surv = compute_survival(torch.rand(12, BLOCK)) + lens = schedule_verify_lens_topk(survival=surv, budget=20, cfg=cfg) + assert lens.dtype == torch.int32 + assert torch.all(lens >= cfg.min_verify_len) + assert torch.all(lens <= cfg.resolved_max_verify_len) + # Budget is spent above the floor and never exceeded. + assert int((lens - cfg.min_verify_len).sum()) <= 20 + # Zero budget degenerates to the floor for everyone. + cfg_floor = _cfg(min_verify_len=2) + lens = schedule_verify_lens_topk( + survival=compute_survival(torch.rand(5, BLOCK)), budget=0, cfg=cfg_floor) + assert torch.equal(lens, torch.full((5,), 2, dtype=torch.int32)) + # A huge budget saturates at the cap. + cfg_cap = _cfg(min_verify_len=1, max_verify_len=4) + lens = schedule_verify_lens_topk( + survival=compute_survival(torch.full((5, BLOCK), 0.99)), budget=10**6, cfg=cfg_cap) + assert torch.equal(lens, torch.full((5,), 4, dtype=torch.int32)) + # An empty batch allocates nothing. + assert schedule_verify_lens_topk( + survival=torch.zeros(0, BLOCK), budget=5, cfg=_cfg()).numel() == 0 + + +def test_survival_eps_excludes_hopeless_positions(): + """Even with unlimited budget, sub-eps candidates are never admitted.""" + cfg = _cfg(min_verify_len=1, survival_eps=1e-3) + conf = torch.tensor([[1.0, 1.0, 1e-9, 1.0, 1.0, 1.0, 1.0]]) + lens = schedule_verify_lens_topk(survival=compute_survival(conf), budget=10**6, cfg=cfg) + # positions 0,1 survive; from position 2 on, survival collapses below eps. + assert int(lens[0]) == 2 + + +def test_allocation_is_deterministic_under_tied_confidences(): + """Two ranks with identical input must produce identical verify lengths.""" + cfg = _cfg(min_verify_len=1) + surv = compute_survival(torch.full((9, BLOCK), 0.9)) + first = schedule_verify_lens_topk(survival=surv, budget=13, cfg=cfg) + for _ in range(8): + assert torch.equal(schedule_verify_lens_topk(survival=surv, budget=13, cfg=cfg), first) + + +def test_ties_break_toward_earlier_positions(): + """All candidates tie, so the allocation must stay a balanced prefix front, + not hand one request a deep suffix while another gets nothing.""" + cfg = _cfg(min_verify_len=1) + surv = compute_survival(torch.full((4, BLOCK), 1.0)) + lens = schedule_verify_lens_topk(survival=surv, budget=4, cfg=cfg) + assert int(lens.max()) - int(lens.min()) <= 1 + + +def test_schedule_config_rejects_bad_bounds(): + with pytest.raises(ValueError, match="min_verify_len must be >= 1"): + DSparkScheduleConfig(block_size=BLOCK, min_verify_len=0) + with pytest.raises(ValueError, match="block_size must be >= 1"): + DSparkScheduleConfig(block_size=0) + with pytest.raises(ValueError, match="< min_verify_len"): + DSparkScheduleConfig(block_size=BLOCK, min_verify_len=5, max_verify_len=3) + with pytest.raises(ValueError, match=r"\[bs, K\]"): + compute_survival(torch.rand(BLOCK)) + + +# -------------------------------------------------------------------------- +# budget conservation: allocation must never exceed what the planner sized +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bs,min_len", list(itertools.product([1, 3, 17], [1, 2]))) +def test_planner_and_allocator_agree_on_the_budget(bs, min_len): + rng = np.random.default_rng(bs * 100 + min_len) + table = SpsCostTable(token_counts=(0, 16, 32, 64, 128), step_time_ms=(3.0, 3.1, 5.0, 5.1, 9.0)) + conf = rng.uniform(0.4, 0.99, size=(bs, BLOCK)) + surv_np = np.cumprod(conf, axis=1) + budget = compute_verify_token_budget( + survival=surv_np, num_gen_requests=bs, cost_table=table, min_verify_len=min_len + ) + lens = schedule_verify_lens_topk( + survival=torch.from_numpy(surv_np).float(), + budget=budget, + cfg=_cfg(min_verify_len=min_len), + ) + assert int((lens - min_len).sum()) <= budget + assert torch.all(lens >= min_len) + + +# -------------------------------------------------------------------------- +# uniform-K projection (what TRT-LLM's graph key can actually express today) +# -------------------------------------------------------------------------- + + +def test_discrete_argmax_beats_round_then_snap_across_a_riser(): + """Optimizing over runnable lengths is not the same as rounding the budget: + snapping down past a cost riser lands on a worse Theta than the search.""" + table = SpsCostTable(token_counts=(0, 8, 24, 40), step_time_ms=(2.0, 2.05, 2.1, 20.0)) + surv = np.cumprod(np.full((8, BLOCK), 0.97), axis=1) + chosen = budget_argmax_over_uniform_lens( + survival=surv, + num_gen_requests=8, + cost_table=table, + allowed_lens=[1, 3, 5, 7], + ) + # 8*(7+1)=64 tokens lands on the 20ms riser; the search must avoid it. + assert chosen < 7 + + def theta(length): + # Mirrors the implementation exactly, including the bonus token every + # request submits alongside its drafts. + tau = 8 + float(surv[:, :length].sum()) + return tau / table.step_time(8 * (length + 1), 8) + + assert theta(chosen) == max(theta(v) for v in (1, 3, 5, 7)) + # With nothing runnable, fall back to the floor. + assert ( + budget_argmax_over_uniform_lens( + survival=surv, num_gen_requests=8, cost_table=table, allowed_lens=[] + ) + == 1 + ) + + +def test_derived_tiers_lose_nothing_versus_continuous_k(): + """The shelf-right-edge property: within a shelf Theta rises monotonically, + so tiers built from right edges must match an exhaustive length search.""" + rng = np.random.default_rng(11) + # Shelves are encoded as breakpoint pairs with equal values: under the + # interpolating consumer, flatness between points is only real when the + # table measured it. + table = SpsCostTable( + token_counts=(0, 11, 12, 29, 30, 47, 48, 89, 90, 159, 160), + step_time_ms=(2.0, 2.0, 2.4, 2.4, 3.9, 3.9, 4.0, 4.0, 7.5, 7.5, 12.0), + ) + for bs in (3, 8, 16, 31): + surv = np.cumprod(rng.uniform(0.5, 0.995, size=(bs, BLOCK)), axis=1) + tiers = derive_verify_len_tiers( + cost_table=table, num_requests=bs, block_size=BLOCK, max_tiers=99 + ) + + def theta(length): + return (bs + float(surv[:, :length].sum())) / table.step_time(bs * (length + 1), bs) + + best_any = max(theta(v) for v in range(1, BLOCK + 1)) + best_tier = max(theta(v) for v in tiers) + assert best_tier == pytest.approx(best_any), f"bs={bs} tiers={tiers}" + + +def test_derived_tiers_respect_the_capture_budget(): + table = SpsCostTable( + token_counts=(0, 5, 9, 14, 20, 27, 35), step_time_ms=(1.0, 2, 3, 4, 5, 6, 7) + ) + tiers = derive_verify_len_tiers(cost_table=table, num_requests=2, block_size=BLOCK, max_tiers=3) + assert len(tiers) <= 3 + # Endpoints are always kept: the floor is always runnable and the full block + # is always a right edge (the last shelf is unbounded). + assert tiers[0] == 1 and tiers[-1] == BLOCK + + +def test_derived_tiers_are_a_function_of_batch_size(): + """Tiers move with bs because total tokens = bs * (length + 1); a tier set + derived once and reused across batch sizes would be wrong for one of them.""" + table = SpsCostTable(token_counts=(0, 40), step_time_ms=(1.0, 5.0)) + assert derive_verify_len_tiers( + cost_table=table, num_requests=8, block_size=BLOCK, max_tiers=99 + ) == [1, 3, BLOCK] + assert derive_verify_len_tiers( + cost_table=table, num_requests=32, block_size=BLOCK, max_tiers=99 + ) == [1, BLOCK] + + +# -------------------------------------------------------------------------- +# host-side planner: relay, fallbacks, cross-rank agreement +# -------------------------------------------------------------------------- + + +def _planner(**kw): + kw.setdefault("cfg", _cfg(min_verify_len=1)) + kw.setdefault( + "cost_table", + SpsCostTable(token_counts=(0, 8, 24, 48), step_time_ms=(2.0, 2.1, 4.0, 9.0)), + ) + kw.setdefault("tiers", [1, 3, BLOCK]) + return DSparkVerifyPlanner(**kw) + + +def test_planner_reads_confidence_by_row_not_by_batch_position(): + """The snapshot is slot-indexed and read one iteration later; ``rows`` is + what re-associates it after joins and departures reshuffle the batch.""" + p = _planner() + # Row 3 is a confident request, row 0 a hopeless one. + buf = torch.full((5, BLOCK), -8.0) + buf[3] = 8.0 + p._host_buffer, p._copy_event, p._snapshot_valid = buf, None, True + + # This step schedules the confident request first, then the hopeless one. + lens = p.decide_verify_lens(num_gen_requests=2, rows=[3, 0]) + assert lens is not None and lens[0] > lens[1], f"row mapping ignored: {lens}" + + +def test_planner_rejects_a_row_list_that_does_not_match_the_batch(): + p = _planner() + p._host_buffer, p._copy_event, p._snapshot_valid = torch.rand(8, BLOCK), None, True + assert p.decide_verify_lens(num_gen_requests=4, rows=[0, 1]) is None + assert p.stats["fallback_short_snapshot"] == 1 + + +def _capture_set(batch_sizes, tiers, max_draft_len=BLOCK): + """Invoke PyTorchModelEngine._get_graphs_to_capture with a stub engine.""" + import types + + from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + + spec_config = types.SimpleNamespace( + enable_confidence_scheduling=True, + verify_len_tiers=list(tiers), + draft_len_schedule=None, + max_concurrency=None, + is_linear_tree=True, + spec_dec_mode=types.SimpleNamespace( + support_dynamic_draft_len=lambda: False, use_one_engine=lambda: False + ), + ) + engine = types.SimpleNamespace( + is_draft_model=False, + spec_config=spec_config, + max_draft_len=max_draft_len, + max_total_draft_tokens=max_draft_len, + original_max_draft_len=max_draft_len, + _dynamic_draft_len_mapping=None, + _cuda_graph_batch_sizes=list(batch_sizes), + ) + return PyTorchModelEngine._get_graphs_to_capture(engine, list(batch_sizes), None) + + +def test_capture_pins_the_draft_length_to_the_top_tier(): + """One graph per batch size, not a cross product with the tier ladder: + nothing varies draft_len at runtime, so lower-tier graphs are unreachable.""" + bss, tiers = [1, 8, 64], [1, 3, BLOCK] + graphs = _capture_set(bss, tiers) + assert set(graphs) == {(bs, BLOCK) for bs in bss} + assert len(graphs) == len(bss), ( + f"expected one graph per batch size, got {len(graphs)}: {graphs}") + + +# -------------------------------------------------------------------------- +# top-k allocation follows confidence +# -------------------------------------------------------------------------- + + +def _block5_cfg() -> DSparkScheduleConfig: + return DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + + +def test_topk_hands_longer_windows_to_more_confident_requests(): + """The budget must follow survival, not batch position.""" + # survival[r, k] = P(the first k+1 drafted tokens all get accepted) + survival = torch.tensor([ + [0.99, 0.98, 0.97, 0.96, 0.95], # confident + [0.90, 0.70, 0.50, 0.30, 0.10], # middling + [0.20, 0.04, 0.01, 0.00, 0.00], # collapses + ]) + lens = schedule_verify_lens_topk(survival=survival, budget=6, + cfg=_block5_cfg()).tolist() + + assert len(lens) == 3 + assert lens[0] > lens[2], ( + f"confident request got {lens[0]}, collapsing request got {lens[2]}; " + f"the budget must follow survival, not batch position") + assert lens[0] >= lens[1] >= lens[2] + # Every request keeps at least the floor, and the budget is respected. + assert min(lens) >= 1 + assert sum(lens) - len(lens) * 1 <= 6 + + +def test_a_full_budget_degenerates_to_the_uniform_full_window(): + """The no-trim case must stay exactly uniform: ragged must cost nothing.""" + num_reqs, max_len = 4, 5 + survival = torch.full((num_reqs, max_len), 0.99) + budget = num_reqs * (max_len - 1) + lens = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=_block5_cfg()).tolist() + assert lens == [max_len] * num_reqs + + +# -------------------------------------------------------------------------- +# cost-table interpolation: clamped linear on both the token and batch axes +# -------------------------------------------------------------------------- + +# The certified GB300 table's shape, abbreviated to two breakpoints. +TABLE = SpsCostTable(token_counts=(512, 768, 1536), + step_time_ms=(23.85, 35.61, 105.64), + fixed_overhead_ms=25.244, + batch_sizes=(128, 256), + batch_overhead_ms=(11.462, 19.343)) + + +def test_theta_interpolates_between_breakpoints(): + """Exact on breakpoints, linear between them: 1512 tokens must cost ~the + 1536 price, not the 768 one a floor lookup would return.""" + assert TABLE.step_time(768, 256) == pytest.approx(25.244 + 19.343 + 35.61) + assert TABLE.step_time(1536, 256) == pytest.approx(25.244 + 19.343 + 105.64) + got = TABLE.step_time(1512, 252) + theta = 35.61 + (105.64 - 35.61) * (1512 - 768) / (1536 - 768) + alpha = 11.462 + (19.343 - 11.462) * (252 - 128) / (256 - 128) + assert got == pytest.approx(25.244 + alpha + theta) + # The floor price it used to return -- and must never return again. + assert got > 130.0 + + +def test_alpha_interpolates_and_clamps(): + assert TABLE.batch_overhead(128) == pytest.approx(11.462) + assert TABLE.batch_overhead(192) == pytest.approx((11.462 + 19.343) / 2) + assert TABLE.batch_overhead(64) == pytest.approx(11.462) # clamp low + assert TABLE.batch_overhead(512) == pytest.approx(19.343) # clamp high + + +def test_theta_clamps_outside_the_measured_range(): + """Theta clamps to the end values; a table without a batch axis adds no alpha.""" + plain = SpsCostTable(token_counts=(512, 768, 1536), + step_time_ms=(23.85, 35.61, 105.64), + fixed_overhead_ms=25.244) + assert plain.step_time(100, 0) == pytest.approx(25.244 + 23.85) + assert plain.step_time(4096, 0) == pytest.approx(25.244 + 105.64) + + +# -------------------------------------------------------------------------- +# decision-level parity with the SGLang budget planner this was ported from +# -------------------------------------------------------------------------- + +# Reference implementation, copied from SGLang dspark_components/dspark_planner +# (additive-table path) with only cosmetic renames. Do not "fix" this copy: +# its job is to be what SGLang runs. + + +def _sgl_interp_clamped(xs, ys, x: float) -> float: + xs = torch.tensor(xs, dtype=torch.float64) + ys = torch.tensor(ys, dtype=torch.float64) + x_t = torch.tensor(float(x), dtype=torch.float64).clamp_(xs[0], xs[-1]) + hi = torch.bucketize(x_t, xs, right=True).clamp_(1, xs.numel() - 1) + lo = hi - 1 + span = (xs[hi] - xs[lo]).clamp_(min=1e-9) + frac = (x_t - xs[lo]) / span + return float(ys[lo] + frac * (ys[hi] - ys[lo])) + + +def _sgl_additive_step_time(table, num_requests: int, num_budgets: int): + floor = table["bias_seconds"] + _sgl_interp_clamped( + table["bs_probes"], table["alpha_seconds"], float(num_requests)) + m_probes = torch.tensor(table["m_probes"], dtype=torch.float64) + theta_vals = torch.tensor(table["theta_seconds"], dtype=torch.float64) + m = (num_requests + torch.arange(num_budgets, dtype=torch.float64)).clamp_( + min=float(table["m_probes"][0]), max=float(table["m_probes"][-1])) + hi = torch.bucketize(m, m_probes, right=True).clamp_(1, m_probes.numel() - 1) + lo = hi - 1 + span = (m_probes[hi] - m_probes[lo]).clamp_(min=1e-9) + frac = (m - m_probes[lo]) / span + theta_at_m = theta_vals[lo] + frac * (theta_vals[hi] - theta_vals[lo]) + return floor + theta_at_m + + +def _sgl_compute_verify_token_budget(*, history_survival_probs, table, + max_verify_len, survival_eps): + num_requests = history_survival_probs.shape[0] + candidates = history_survival_probs[:, :max_verify_len].flatten() + candidates = candidates[candidates >= survival_eps].to(torch.float64) + candidates_sorted = torch.sort(candidates, descending=True).values + prefix_sum = torch.cumsum(candidates_sorted, dim=0) + tau_star = num_requests + torch.cat( + [torch.zeros(1, dtype=torch.float64), prefix_sum]) + step_time = _sgl_additive_step_time(table, int(num_requests), + int(tau_star.numel())) + theta = tau_star / step_time + return int(torch.argmax(theta)) + + +# One certified-table shape shared by both sides. TensorRT-LLM stores ms, +# SGLang stores seconds; the argmax is scale-invariant but the tables are +# built with the honest factor anyway. + +M_PROBES = (64, 96, 128, 192, 384, 512, 768, 1536) +THETA_MS = (9.68, 11.08, 11.95, 13.95, 22.09, 23.85, 35.61, 105.64) +BIAS_MS = 25.244443 +BS_PROBES = (32, 64, 128, 256) +ALPHA_MS = (0.0, 2.177477, 11.462212, 19.342661) +SGL_BLOCK = 5 + +TRT_TABLE = SpsCostTable(token_counts=M_PROBES, step_time_ms=THETA_MS, + fixed_overhead_ms=BIAS_MS, batch_sizes=BS_PROBES, + batch_overhead_ms=ALPHA_MS) +SGL_TABLE = { + "m_probes": M_PROBES, + "theta_seconds": tuple(v / 1e3 for v in THETA_MS), + "bias_seconds": BIAS_MS / 1e3, + "bs_probes": BS_PROBES, + "alpha_seconds": tuple(v / 1e3 for v in ALPHA_MS), +} + + +def _inversion_free_survival(rng, bs): + """Survivals where every position-0 beats every deeper candidate; the + shift-exact parity theorem holds only under this condition.""" + conf0 = rng.uniform(0.85, 0.99, size=(bs, 1)) + deeper = rng.uniform(0.25, 0.80, size=(bs, SGL_BLOCK - 1)) + return np.cumprod(np.concatenate([conf0, deeper], axis=1), axis=1) + + +@pytest.mark.parametrize("bs", [1, 3, 8, 64, 128, 252, 256]) +def test_budget_argmax_matches_sglang_up_to_the_floor_shift(bs): + """Same survival, same table, eps off, no inversions: shift-exact parity.""" + rng = np.random.default_rng(17 + bs) + for _ in range(5): + surv = _inversion_free_survival(rng, bs) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=SGL_BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + if sgl_n >= bs: + assert trt_m == sgl_n - bs, ( + f"bs={bs}: SGLang admits {sgl_n} candidates " + f"(= floor {bs} + {sgl_n - bs}), TRT-LLM budget {trt_m}") + else: + # SGLang went below the one-draft floor; the closest budget + # TensorRT-LLM can express under min_verify_len=1 is zero. + assert trt_m == 0 + + +def test_live_pooled_survivals_park_on_the_breakpoint_together(): + """The campaign's operating point: both planners stop at exactly M = 768, + and tier alignment then rounds down to the capturable rung-2 budget.""" + bs = 252 + surv = np.tile([0.771, 0.693, 0.502, 0.360, 0.315], (bs, 1)) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=SGL_BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + assert trt_m == 768 - 2 * bs + assert sgl_n - bs == trt_m + tiered = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1, allowed_lens=[1, 2, 5]) + assert tiered == bs * 1 # rung-2: one scheduled position past the floor + + +def test_starving_weak_rows_is_a_known_gap(): + """SGLang can verify NOTHING for a hopeless request while TRT-LLM's + min_verify_len=1 floor cannot, so SGLang's achievable Theta is better.""" + # Half the batch is strong all the way down, half is dead on arrival. + # The batch must be big enough that its token range sits on a genuinely + # rising part of the theta curve, or every candidate is free and the + # comparison proves nothing. + bs = 64 + strong = np.tile([0.99, 0.98, 0.97, 0.96, 0.95], (bs // 2, 1)) + dead = np.tile([0.01, 0.008, 0.006, 0.004, 0.002], (bs // 2, 1)) + surv = np.concatenate([strong, dead], axis=0) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=SGL_BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + + def theta_at(tau, tokens): + return tau / float(TRT_TABLE.step_times(np.asarray([tokens]), bs)[0]) + + cand = np.sort(surv.reshape(-1))[::-1] + sgl_theta = theta_at(bs + cand[:sgl_n].sum(), bs + sgl_n) + deeper = np.sort(surv[:, 1:].reshape(-1))[::-1] + trt_theta = theta_at(bs + surv[:, 0].sum() + deeper[:trt_m].sum(), + 2 * bs + trt_m) + assert sgl_theta > trt_theta + + +def test_eps_filtering_is_a_known_divergence(): + """SGLang drops sub-eps candidates inside the budget; TRT-LLM applies eps + at allocation time, so its budget can count a candidate the allocator refuses.""" + bs = 16 + surv = np.tile([0.9, 0.4, 0.008, 0.004, 0.002], (bs, 1)) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=SGL_BLOCK, survival_eps=0.01) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + assert sgl_n - bs <= trt_m + + +# -------------------------------------------------------------------------- +# tier-aligned budgets: score with the tau the step actually collects +# -------------------------------------------------------------------------- + + +def _table() -> SpsCostTable: + """A staircase with genuine risers, so trimming is worth something.""" + token_counts = tuple(range(0, 400, 16)) + step_time_ms = tuple(4.0 + 0.6 * (tok // 96) + 0.004 * tok + for tok in token_counts) + return SpsCostTable(token_counts=token_counts, step_time_ms=step_time_ms, + fixed_overhead_ms=1.0) + + +@pytest.mark.parametrize("tiers", [[1, 2, 5], [1, 3, 5], [1, 5]]) +def test_restricted_answer_is_always_realisable(tiers): + """Every returned budget corresponds to a rung the executor captured.""" + rng = np.random.default_rng(20260803) + table = _table() + for _ in range(200): + bs = int(rng.integers(2, 17)) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] + n = compute_verify_token_budget( + survival=survival, num_gen_requests=bs, cost_table=table, + min_verify_len=1, allowed_lens=tiers) + assert n % bs == 0, ( + f"budget {n} for {bs} requests is not n*(t-min) for any tier") + assert (n // bs) + 1 in tiers, ( + f"budget {n} implies tier {(n // bs) + 1}, not in {tiers}") + + +def test_restricted_never_scores_worse_than_the_uniform_choice(): + """Same grid, better numerator: the chosen rung's realised theta wins; + at bs=1 the two scorers must pick the same rung outright.""" + rng = np.random.default_rng(31337) + table = _table() + tiers = [1, 2, 5] + for _ in range(300): + bs = int(rng.integers(1, 17)) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] + + def realised_theta(tier: int) -> float: + budget = bs * (tier - 1) + cand = np.sort(survival[:, 1:].reshape(-1))[::-1] + tau = float(bs) + float(survival[:, :1].sum()) + float( + cand[:budget].sum()) + tokens = np.array([bs * (tier + 1)]) + return tau / float(table.step_times(tokens, bs)[0]) + + n = compute_verify_token_budget( + survival=survival, num_gen_requests=bs, cost_table=table, + min_verify_len=1, allowed_lens=tiers) + chosen = (n // bs) + 1 + uniform = budget_argmax_over_uniform_lens( + survival=survival, num_gen_requests=bs, cost_table=table, + allowed_lens=tiers, min_verify_len=1) + assert realised_theta(chosen) >= realised_theta(uniform) - 1e-12, ( + f"restricted picked tier {chosen} whose realised theta is below " + f"the uniform scorer's tier {uniform}") + if bs == 1: + assert chosen == uniform, ( + "at bs=1 the uniform and top-k allocations are the same set; " + "the two scorers must pick the same rung") + + +# -------------------------------------------------------------------------- +# runtime verify-length pin: queue, agree, adopt +# -------------------------------------------------------------------------- + + +def _pin_planner(tiers=(1, 2, 5)): + table = SpsCostTable(token_counts=(0, 512, 768, 1536), + step_time_ms=(5.0, 68.4, 80.2, 150.5), + fixed_overhead_ms=1.0) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + return DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=list(tiers)) + + +def test_adopting_the_agreed_value_applies_it(): + """Queuing must not change this rank's behaviour before the group agrees; + adopting the agreed value applies it and consumes the pending request.""" + planner = _pin_planner() + assert planner.request_verify_len_pin(2) == 2 + # Queued, not yet in force. + assert planner._forced_verify_len is None + assert planner.pending_verify_len_pin() == 2 + planner.adopt_verify_len_pin(2) + assert planner._forced_verify_len == 2 + # Consumed: the next step must not re-broadcast a pin nobody asked for. + assert planner.pending_verify_len_pin() == -1 + + +def test_a_rank_that_queued_nothing_still_adopts_the_group_value(): + """The endpoint lands on ONE rank; the others learn it from the payload.""" + peer = _pin_planner() + assert peer.pending_verify_len_pin() == -1 + peer.adopt_verify_len_pin(5) + assert peer._forced_verify_len == 5 + + +def test_zero_clears_the_pin(): + """Wire protocol: -1 means nobody asked (leave the pin alone); clearing + needs its own value, 0, because None cannot travel in an int payload.""" + planner = _pin_planner() + planner.adopt_verify_len_pin(2) + planner.adopt_verify_len_pin(-1) + assert planner._forced_verify_len == 2 + assert planner.request_verify_len_pin(None) is None + assert planner.pending_verify_len_pin() == 0 + planner.adopt_verify_len_pin(0) + assert planner._forced_verify_len is None + + +def test_an_uncaptured_length_is_refused_at_the_call_site(): + """Rejected when requested, not on some later step, and never queued.""" + planner = _pin_planner(tiers=(1, 2, 5)) + with pytest.raises(ValueError, match="captured tier ladder"): + planner.request_verify_len_pin(3) + with pytest.raises(ValueError, match="outside"): + planner.request_verify_len_pin(9) + assert planner.pending_verify_len_pin() == -1 + assert planner._forced_verify_len is None + + +def test_pinned_steps_hand_every_request_the_same_window(): + """The whole point: the shape a sweep cell can honestly label.""" + planner = _pin_planner() + bs = 8 + rng = np.random.default_rng(5) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + planner.adopt_verify_len_pin(2) + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens == [2] * bs + assert planner.stats["forced_steps"] == 1 + + +# -------------------------------------------------------------------------- +# snapshot instruments: stamp-lag histogram and neutral-row count +# -------------------------------------------------------------------------- + + +def test_neutral_rows_are_counted_and_real_rows_are_not(): + planner = _pin_planner() + selected = torch.zeros((8, 5)) + selected[3] = 30.0 # the neutral fill: sigmoid ~ 1.0 at any temperature + planner._note_snapshot_stats(selected, rows=None) + assert planner.stats["snap_rows"] == 8 + assert planner.stats["snap_neutral_rows"] == 1 + + +def test_lag_histogram_reads_staged_seq_minus_stamp(): + planner = _pin_planner() + planner._host_stamps = torch.tensor([7, 6, 5, 0], dtype=torch.int32) + planner._staged_seq = 7 + selected = torch.zeros((4, 5)) + planner._note_snapshot_stats(selected, rows=[0, 1, 2, 3]) + hist = planner.stats["stamp_lag_hist"] + # Stamp 0 is "never drafted": lag 7 in truth, and it must land visibly at + # the histogram's far end rather than vanish -- these are exactly the rows + # the neutral counter should agree about. + assert hist == {0: 1, 1: 1, 2: 1, 7: 1} + + +def test_large_and_negative_lags_clamp_instead_of_exploding_the_dict(): + planner = _pin_planner() + planner._host_stamps = torch.tensor([0, 500], dtype=torch.int32) + planner._staged_seq = 400 + planner._note_snapshot_stats(torch.zeros((2, 5)), rows=[0, 1]) + assert planner.stats["stamp_lag_hist"] == {8: 1, -2: 1} + + +def test_counts_accumulate_across_steps(): + planner = _pin_planner() + planner._host_stamps = torch.tensor([3, 3], dtype=torch.int32) + planner._staged_seq = 4 + for _ in range(3): + planner._note_snapshot_stats(torch.full((2, 5), 30.0), rows=[0, 1]) + assert planner.stats["snap_rows"] == 6 + assert planner.stats["snap_neutral_rows"] == 6 + assert planner.stats["stamp_lag_hist"] == {1: 6} + + +def test_instrument_failure_is_counted_not_raised(): + planner = _pin_planner() + planner._host_stamps = torch.tensor([1], dtype=torch.int32) + planner._staged_seq = 1 + # Rows that cannot index the stamp buffer must clamp, and anything worse + # must be swallowed into the error counter: instruments never kill steps. + planner._note_snapshot_stats(torch.zeros((1, 5)), rows=[10_000]) + assert planner.stats.get("snap_stats_errors", 0) == 0 + planner._note_snapshot_stats(object(), rows=None) # type: ignore[arg-type] + assert planner.stats["snap_stats_errors"] == 1 + + +def test_decide_path_feeds_the_instruments(): + """The counters must fill from the real decide path, not only direct calls.""" + planner = _pin_planner() + bs = 8 + rng = np.random.default_rng(3) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens is not None + assert planner.stats["snap_rows"] == bs + assert planner.stats["snap_neutral_rows"] == 0 + # The local-decision instruments must fill on the same path: one rung per + # decision, one window per request, and the two must agree -- the windows + # are the top-k realisation of exactly that rung's budget. + rung_hist = planner.stats["local_rung_hist"] + assert sum(rung_hist.values()) == 1 + (rung,) = rung_hist + assert rung in (1, 2, 5) + len_hist = planner.stats["local_len_hist"] + assert sum(len_hist.values()) == bs + budget = bs * (rung - planner.cfg.min_verify_len) + extra = sum((k - planner.cfg.min_verify_len) * v for k, v in len_hist.items()) + assert extra <= budget + + +# -------------------------------------------------------------------------- +# reconciliation guards: engine fingerprint and price prediction +# -------------------------------------------------------------------------- + +LIVE = {"tp": 8, "ep": 8, "attention_dp": True, "block": 5, + "max_batch_size": 256} + + +def _payload(engine): + return {"token_counts": [512, 1536], "step_time_ms": [60.0, 150.0], + "_meta": {"engine": engine} if engine is not None else {}} + + +def test_mismatched_engine_is_refused(): + """Same cell, different max_batch_size: the wrong table must not load.""" + wrong = dict(LIVE, max_batch_size=64) + with pytest.raises(ValueError, match="different engine configuration"): + check_table_fingerprint(payload=_payload(wrong), live=dict(LIVE)) + + +def test_fingerprints_that_must_load(): + """Every non-contradicted fingerprint loads; only a contradicted fact + refuses (case-insensitive strings, unverifiable keys logged, no-meta warns).""" + check_table_fingerprint(payload=_payload(dict(LIVE)), live=dict(LIVE)) + check_table_fingerprint( + payload=_payload(dict(LIVE, moe_backend="megamoe_cutedsl")), + live=dict(LIVE, moe_backend="MEGAMOE_CUTEDSL")) + check_table_fingerprint( + payload=_payload(dict(LIVE, image="faf2c60935", + geometry="constant_block")), + live=dict(LIVE)) + check_table_fingerprint(payload=_payload(None), live=dict(LIVE)) + + +def test_planner_records_the_price_it_paid(): + """Every budget decision must leave behind the step time it assumed, so a + gap against measured hostStepTimeMS can contradict a wrong table.""" + table = SpsCostTable(token_counts=(0, 512, 768, 1536), + step_time_ms=(5.0, 68.4, 80.2, 150.5), + fixed_overhead_ms=1.0) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + planner = DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=[1, 2, 5]) + + bs = 8 + rng = np.random.default_rng(11) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + + # Only the snapshot is faked; calibration stays the constructor's default + # sigmoid (setting it to None post-construction would bypass the + # `apply_calibration or torch.sigmoid` fallback and crash the decide path). + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens is not None + + steps = planner.stats.get("predicted_steps", 0) + total = planner.stats.get("predicted_ms_sum", 0.0) + assert steps == 1 and total > 0.0 + + # The recorded price must be the table's own answer for the tokens the + # decision implies -- anything else and the reconciliation compares two + # unrelated numbers. + budget = sum(lens) - bs * cfg.min_verify_len + tokens = bs * (cfg.min_verify_len + 1) + budget + expected = float(table.step_times(np.asarray([tokens]), bs)[0]) + assert planner.last_predicted_step_ms == pytest.approx(expected) + assert total == pytest.approx(expected) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py new file mode 100644 index 000000000000..15f4885cd679 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py @@ -0,0 +1,1199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DSpark ragged (per-request) verification subsystem: packed layout and +bucket selection, host/device fill parity, overlap-scheduler gather layout, +compressor per-request new-token counts, and the ragged CUDA-graph key.""" + +import types + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import \ + DeepseekV4TrtllmAttentionMetadata +from tensorrt_llm._torch.pyexecutor.llm_request import get_request_tokens_per_gen_step +from tensorrt_llm._torch.speculative.dspark_device_select import \ + select_windows_device +from tensorrt_llm._torch.speculative.dspark_ragged import ( + RaggedVerifyLayout, + build_qo_indptr, + build_row_maps_device, + choose_ragged_capture_shape, + count_accepted_ragged, + fill_bucket_device, + fill_padded_rows_onehot, + ragged_gather_index_lists, + round_up_to_bucket, + row_ids_from_lens, + scatter_ragged_to_padded, +) +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, + compute_survival, + schedule_verify_lens_topk, +) + +BLOCK = 7 + + +def _layout(lens, bucket=None, buckets=None): + t = torch.tensor(lens, dtype=torch.int32) + return RaggedVerifyLayout.from_verify_lens( + t, graph_num_tokens=bucket, buckets=buckets, total_verify_tokens=sum(lens) + ) + + +# -------------------------------------------------------------------------- +# packing +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "lens, expected", + [ + ([3, 1, 4, 2], [0, 3, 4, 8, 10]), + ([], [0]), # empty batch still yields the leading 0 + ], +) +def test_qo_indptr_is_the_exclusive_prefix_sum(lens, expected): + assert build_qo_indptr(torch.tensor(lens, dtype=torch.int32)).tolist() == expected + + +def test_build_qo_indptr_rejects_2d(): + with pytest.raises(ValueError, match="1-D"): + build_qo_indptr(torch.ones(2, 3, dtype=torch.int32)) + + +# -------------------------------------------------------------------------- +# bucket selection +# -------------------------------------------------------------------------- + + +def test_round_up_picks_the_smallest_fitting_bucket(): + b = [8, 16, 32, 64] + assert round_up_to_bucket(1, b) == 8 + assert round_up_to_bucket(8, b) == 8 + assert round_up_to_bucket(9, b) == 16 + assert round_up_to_bucket(64, b) == 64 + + +def test_round_up_raises_past_the_largest_bucket(): + """Clamping would silently drop the step out of graph replay into eager.""" + with pytest.raises(ValueError, match="exceeds the largest captured bucket"): + round_up_to_bucket(65, [8, 16, 32, 64]) + + +def test_layout_derives_the_bucket_from_the_host_total(): + lay = RaggedVerifyLayout.from_verify_lens( + torch.tensor([3, 1, 4], dtype=torch.int32), buckets=[8, 16, 32], total_verify_tokens=8 + ) + assert lay.graph_num_tokens == 8 + # extend_start_loc is the offset form some attention backends want directly. + assert torch.equal(lay.extend_start_loc, lay.qo_indptr[:-1]) + + +def test_layout_refuses_to_sync_for_the_total(): + with pytest.raises(ValueError, match="would sync"): + RaggedVerifyLayout.from_verify_lens(torch.tensor([3, 1], dtype=torch.int32)) + + +# -------------------------------------------------------------------------- +# padding reclamation +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "lens, bucket", + [ + ([2, 1, 1], 12), # 4 used, 8 spare: padding becomes real verification + ([3, 2, 3], 8), # already exact: the fill is a no-op + ], +) +def test_fill_bucket_converts_padding_into_real_verification(lens, bucket): + """The step already pays for the bucket -- spare slots must not be wasted.""" + lay = _layout(lens, bucket=bucket) + filled = lay.fill_bucket(max_verify_len=BLOCK) + assert filled.total_verify_tokens == bucket + assert int(filled.verify_lens.sum()) == bucket + filled.validate(max_verify_len=BLOCK, exact_fill=True) + + +def test_fill_bucket_spreads_over_real_requests_first(): + """Round-robin keeps the extra where survival is still plausible.""" + lay = _layout([1, 1, 1, 1], bucket=8) + filled = lay.fill_bucket(max_verify_len=BLOCK) + lens = filled.verify_lens.tolist() + assert sum(lens) == 8 + assert max(lens) - min(lens) <= 1, f"unbalanced: {lens}" + + +def test_fill_bucket_adds_pad_rows_to_reach_the_captured_batch_size(): + """Row count must hit padded_bs; slack goes to real rows before pad rows.""" + lay = _layout([3, 2], bucket=12) + filled = lay.fill_bucket(max_verify_len=BLOCK, padded_bs=4) + assert filled.bs == 4 + assert filled.num_real_requests == 2 + assert filled.num_pad_requests == 2 + assert int(filled.verify_lens.sum()) == 12 + filled.validate(max_verify_len=BLOCK, exact_fill=True) + # Real rows absorb all the slack; the pad rows stay at the 1-token minimum + # (a pad row with 0 tokens would break qo_indptr slicing). + lens = filled.verify_lens.tolist() + assert lens[2:] == [1, 1] + assert lens[0] + lens[1] == 10 + + +def test_fill_bucket_raises_when_the_bucket_cannot_be_reached(): + """Silently returning a short batch would desync attention from the MoE.""" + lay = _layout([1, 1], bucket=100) # capacity 2*7 = 14 < 100 + with pytest.raises(ValueError, match="exceeds what .* rows can absorb"): + lay.fill_bucket(max_verify_len=BLOCK) + + +@pytest.mark.parametrize( + "lens, bucket, padded_bs", + [ + ([5, 5], 8, None), # 10 real tokens already > 8 + # 8 real tokens + 6 one-token pad rows needs 14, not 10: the pad-row + # floor counts against the bucket too. A short batch would desync + # attention from the MoE. + ([4, 4], 10, 8), + ], +) +def test_fill_bucket_raises_when_the_bucket_is_too_small(lens, bucket, padded_bs): + lay = _layout(lens, bucket=bucket) + with pytest.raises(ValueError, match="too small"): + lay.fill_bucket(max_verify_len=BLOCK, padded_bs=padded_bs) + + +def test_fill_bucket_rejects_a_padded_bs_below_the_real_count(): + lay = _layout([2, 2, 2], bucket=8) + with pytest.raises(ValueError, match="cannot hold this batch"): + lay.fill_bucket(max_verify_len=BLOCK, padded_bs=2) + + +def test_validate_exact_fill_catches_a_short_batch(): + """The check that would otherwise only surface as an attention/MoE mismatch.""" + lay = _layout([2, 2], bucket=8) # 4 != 8 + lay.validate(max_verify_len=BLOCK) # non-exact form still passes + with pytest.raises(ValueError, match="do not fill the captured bucket"): + lay.validate(max_verify_len=BLOCK, exact_fill=True) + + +def test_fill_bucket_needs_the_host_total(): + lay = RaggedVerifyLayout.from_verify_lens( + torch.tensor([2, 2], dtype=torch.int32), graph_num_tokens=8 + ) + with pytest.raises(ValueError, match="host-side total"): + lay.fill_bucket(max_verify_len=BLOCK) + + +# -------------------------------------------------------------------------- +# invariants +# -------------------------------------------------------------------------- + + +def test_validate_rejects_a_request_that_verifies_nothing(): + """verify_len 0 means the request makes no progress -- it would stall.""" + lay = _layout([2, 0, 3], bucket=8) + with pytest.raises(ValueError, match="anchor token"): + lay.validate() + + +def test_validate_rejects_overflowing_the_bucket(): + lay = _layout([5, 5], bucket=8) + with pytest.raises(ValueError, match="exceed the captured bucket"): + lay.validate() + + +def test_validate_rejects_a_stale_indptr(): + lay = _layout([2, 3], bucket=8) + lay.qo_indptr = torch.tensor([0, 1, 5], dtype=torch.int32) + with pytest.raises(ValueError, match="not the prefix sum"): + lay.validate() + + +# -------------------------------------------------------------------------- +# end-to-end: confidence -> per-request lengths -> packed layout +# -------------------------------------------------------------------------- + + +def test_ragged_beats_uniform_at_equal_budget(): + """Heterogeneous confidence must yield heterogeneous lengths that pack into + a valid layout and out-earn the best uniform split at the same token spend.""" + conf = torch.tensor( + [ + [0.98, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92], # deep, worth verifying + [0.55, 0.40, 0.30, 0.25, 0.20, 0.15, 0.10], # shallow + [0.90, 0.85, 0.70, 0.55, 0.45, 0.35, 0.25], + [0.60, 0.45, 0.35, 0.28, 0.22, 0.18, 0.12], + ] + ) + surv = compute_survival(conf) + cfg = DSparkScheduleConfig(block_size=BLOCK, min_verify_len=1) + budget, bs = 8, conf.shape[0] + + lens = schedule_verify_lens_topk(survival=surv, budget=budget, cfg=cfg) + assert lens[0] > lens[1], "the confident request should verify deeper" + + # The lengths pack into a valid captured layout. + total = int(lens.sum()) + lay = RaggedVerifyLayout.from_verify_lens( + lens, buckets=[8, 16, 32, 64], total_verify_tokens=total + ) + lay.validate(max_verify_len=BLOCK) + assert lay.graph_num_tokens >= total + + # Same tokens verified, more expected yield than the best uniform length. + tau_ragged = sum(float(surv[r, : int(lens[r])].sum()) for r in range(bs)) + budget_tokens = bs * cfg.min_verify_len + budget + tau_uniform = max( + float(surv[:, :L].sum()) for L in range(1, BLOCK + 1) if bs * L <= budget_tokens + ) + assert int(lens.sum()) <= budget_tokens + assert tau_ragged > tau_uniform, ( + f"ragged {tau_ragged:.2f} should beat uniform {tau_uniform:.2f} at the same token spend" + ) + + +# -------------------------------------------------------------------------- +# packed <-> padded, and ragged acceptance counting +# -------------------------------------------------------------------------- + + +def test_row_ids_map_each_packed_token_to_its_request(): + lens = torch.tensor([3, 1, 2], dtype=torch.int32) + assert row_ids_from_lens(lens, total=6).tolist() == [0, 0, 0, 1, 2, 2] + + +def test_row_ids_requires_the_host_total(): + """``total`` keeps the expansion sync-free -- without it torch reads the + cumsum back to the host, which is illegal inside the captured graph.""" + lens = torch.tensor([3, 1, 2], dtype=torch.int32) + with pytest.raises(TypeError): + row_ids_from_lens(lens) + + +def test_scatter_unpacks_a_flat_batch_into_rows(): + lens = torch.tensor([3, 1, 2], dtype=torch.int32) + lay = _layout([3, 1, 2], bucket=8) + flat = torch.tensor([10, 11, 12, 20, 30, 31]) + out = scatter_ragged_to_padded( + flat, verify_lens=lens, qo_indptr=lay.qo_indptr, max_len=4, pad_value=-1 + ) + assert out.tolist() == [[10, 11, 12, -1], [20, -1, -1, -1], [30, 31, -1, -1]] + + +@pytest.mark.parametrize( + "draft, target, lens, expected", + [ + # the run of matches stops at the first mismatch + ([[1, 2, 3, 4]], [[1, 2, 9, 4]], [4], [2]), + # a request only gets credit inside its window, even when the stale + # padding beyond it happens to match -- unmasked padding would credit + # an acceptance the target never made (wrong output, no error). + ([[1, 2, 3, 4], [5, 6, 7, 8]], [[1, 2, 3, 4], [5, 6, 7, 8]], [4, 2], [4, 2]), + ], +) +def test_accept_count_respects_each_request_window(draft, target, lens, expected): + got = count_accepted_ragged( + draft_tokens=torch.tensor(draft), + target_tokens=torch.tensor(target), + verify_lens=torch.tensor(lens), + ) + assert got.tolist() == expected + + +def test_accept_count_matches_the_uniform_path_when_lengths_are_equal(): + """Equal lengths must reproduce the existing rectangular cumprod exactly.""" + torch.manual_seed(9) + bs, L = 6, 5 + draft = torch.randint(0, 4, (bs, L)) + target = torch.randint(0, 4, (bs, L)) + lens = torch.full((bs,), L) + ragged = count_accepted_ragged(draft_tokens=draft, target_tokens=target, verify_lens=lens) + uniform = torch.cumprod((draft == target).int(), dim=-1).sum(1) + assert torch.equal(ragged, uniform) + + +def test_accept_count_rejects_mismatched_shapes(): + with pytest.raises(ValueError, match="same padded shape"): + count_accepted_ragged( + draft_tokens=torch.zeros(2, 3, dtype=torch.long), + target_tokens=torch.zeros(2, 4, dtype=torch.long), + verify_lens=torch.tensor([3, 3]), + ) + + +# --------------------------------------------------------------------------- +# choose_ragged_capture_shape +# --------------------------------------------------------------------------- + +BS_BUCKETS = [1, 2, 4, 8, 16] +TOK_BUCKETS = [8, 16, 32, 64, 128] + + +def _shape(num_real, total, peers=None): + return choose_ragged_capture_shape( + num_real_requests=num_real, + total_verify_tokens=total, + bs_buckets=BS_BUCKETS, + token_buckets=TOK_BUCKETS, + peer_stats=peers, + ) + + +@pytest.mark.parametrize( + "num_real, total, want_bs, want_bucket", + [ + # 5 -> 8 rows; slack = 20 - 5 = 15; needed = 8 + 15 = 23 -> 32. + (5, 20, 8, 32), + # exact captured values are not rounded further: 8 requests with + # slack 8 -> needed 16, both axes already captured. + (8, 16, 8, 16), + ], +) +def test_shape_rounds_both_axes_up_to_captured_values(num_real, total, want_bs, want_bucket): + got = _shape(num_real=num_real, total=total) + assert got.padded_bs == want_bs + assert got.bucket == want_bucket + + +def test_every_rank_computes_the_same_shape(): + # The ADP invariant: ranks with different batches must agree. Each rank + # passes the same peer list and must get an identical answer. + peers = [(3, 9), (6, 30), (5, 11)] + shapes = {_shape(num_real=real, total=total, peers=peers) for real, total in peers} + assert len(shapes) == 1 + + +@pytest.mark.parametrize( + "num_real, total, peers, want_bs, want_bucket", + [ + # This rank is small (3 requests, 9 tokens) but must size for the + # widest peer: rows from max real (6 -> 8), slack from the max + # (30 - 6 = 24) -> 8 + 24 = 32, exactly a captured bucket. + (3, 9, [(3, 9), (6, 30), (5, 11)], 8, 32), + # The widest axes can come from different ranks: rank A has the most + # rows (7 -> 8), rank B the most slack (18) -> 8 + 18 = 26 -> 32. + (7, 8, [(7, 8), (2, 20)], 8, 32), + ], +) +def test_shape_is_driven_by_the_widest_rank_not_the_local_one( + num_real, total, peers, want_bs, want_bucket +): + got = _shape(num_real=num_real, total=total, peers=peers) + assert got.padded_bs == want_bs + assert got.bucket == want_bucket + + +@pytest.mark.parametrize( + "num_real, total, match", + [ + (17, 17, "exceed the largest captured batch"), # rows overflow + (16, 1000, "exceeds the largest captured bucket"), # tokens overflow + (5, 4, "cannot be below"), # fewer tokens than requests + ], +) +def test_shape_raises_when_the_batch_exceeds_captured_rows(num_real, total, match): + with pytest.raises(ValueError, match=match): + _shape(num_real=num_real, total=total) + + +def test_shape_requires_bs_buckets(): + with pytest.raises(ValueError, match="requires bs_buckets"): + choose_ragged_capture_shape( + num_real_requests=1, total_verify_tokens=1, bs_buckets=[], token_buckets=TOK_BUCKETS + ) + + +# --------------------------------------------------------------------------- +# group-consistent bucket choice (max_verify_len given): pad rows share one +# window, so decomposability is a divisibility question every rank must answer +# identically -- a lone decline used to take the whole ADP group eager. +# --------------------------------------------------------------------------- + +MAXLEN_BUCKETS = [16, 24, 32, 40, 48] # padded_bs 8, tiers 1..5 -> 8*(t+1) + + +def _shape6(num_real, total, peers): + return choose_ragged_capture_shape( + num_real_requests=num_real, + total_verify_tokens=total, + bs_buckets=BS_BUCKETS, + token_buckets=MAXLEN_BUCKETS, + peer_stats=peers, + max_verify_len=6, + ) + + +def test_pinned_ranks_agree_on_a_bucket_every_rank_can_decompose(): + """Full-block-pinned ranks must land together on a bucket every rank can + decompose (48 here), not split the group at the first fitting bucket (40).""" + peers = [(5, 30), (4, 24)] + shapes = {_shape6(num_real=r, total=t, peers=peers) for r, t in peers} + assert len(shapes) == 1 + assert shapes.pop().bucket == 48 + + +def test_slacky_batch_keeps_the_smallest_bucket(): + """Feasibility filtering must not grow the bucket when real rows have slack + to absorb the remainder -- the trimmed batch keeps the ungated answer.""" + peers = [(3, 9), (6, 30), (5, 11)] + got = _shape6(num_real=3, total=9, peers=peers) + ungated = choose_ragged_capture_shape( + num_real_requests=3, total_verify_tokens=9, + bs_buckets=BS_BUCKETS, token_buckets=MAXLEN_BUCKETS, + peer_stats=peers) + assert got.bucket == ungated.bucket == 32 + assert got.padded_bs == ungated.padded_bs == 8 + + +def test_group_declines_together_when_nothing_decomposes(): + """When no captured bucket works for every rank, every rank must see the + same ValueError -- a group-consistent decline replays the uniform graph.""" + truncated = [16, 24, 32, 40] # top tier 48 not captured + peers = [(5, 30), (4, 24)] # pinned: 40 fails (5,30), 32 fails it too + for r, t in peers: + with pytest.raises(ValueError, match="decomposable by every rank"): + choose_ragged_capture_shape( + num_real_requests=r, total_verify_tokens=t, + bs_buckets=BS_BUCKETS, token_buckets=truncated, + peer_stats=peers, max_verify_len=6) + + +def test_uniform_batch_shape_matches_the_one_dimensional_rule(): + # With a uniform K the token total is bs * (K + 1), so the joint rule must + # not pick anything larger than the uniform path would have needed. + k_plus_1 = 4 + for num_real in (1, 3, 8): + got = _shape(num_real=num_real, total=num_real * k_plus_1) + assert got.padded_bs >= num_real + assert got.bucket >= num_real * k_plus_1 + + +# --------------------------------------------------------------------------- +# fill_padded_rows_onehot +# --------------------------------------------------------------------------- + + +def test_padding_rows_become_a_valid_distribution(): + probs = torch.zeros(2, 3, 5) + probs[0, :2] = torch.tensor([0.2, 0.3, 0.1, 0.2, 0.2]) + probs[1, :1] = torch.tensor([0.5, 0.1, 0.1, 0.2, 0.1]) + fill_padded_rows_onehot(probs, verify_lens=torch.tensor([2, 1])) + + # Real positions untouched. + assert torch.allclose(probs[0, 0], torch.tensor([0.2, 0.3, 0.1, 0.2, 0.2])) + # Padding positions sum to 1 instead of 0. + assert torch.allclose(probs[0, 2].sum(), torch.tensor(1.0)) + assert torch.allclose(probs[1, 1].sum(), torch.tensor(1.0)) + assert torch.allclose(probs[1, 2].sum(), torch.tensor(1.0)) + + # Full-length rows are left entirely alone (guards the >= boundary). + full = torch.full((2, 3, 4), 0.25) + before = full.clone() + fill_padded_rows_onehot(full, verify_lens=torch.tensor([3, 3])) + assert torch.equal(full, before) + + +# --------------------------------------------------------------------------- +# Bucket fitting: the runtime batch must land on a captured token total +# --------------------------------------------------------------------------- + + +def _fit(real_token_lens, bucket, padded_bs, max_verify_len): + """Mirror of ModelEngine.fit_ragged_verify_lens' core.""" + lens = torch.tensor(real_token_lens, dtype=torch.int32) + layout = RaggedVerifyLayout.from_verify_lens( + lens, graph_num_tokens=bucket, total_verify_tokens=sum(real_token_lens) + ) + return layout.fill_bucket(max_verify_len=max_verify_len, padded_bs=padded_bs) + + +def test_fit_never_exceeds_the_per_request_ceiling(): + got = _fit([1, 1, 1], bucket=12, padded_bs=4, max_verify_len=3) + assert max(got.verify_lens.tolist()) <= 3 + + +def test_warmup_split_reaches_any_captured_bucket(): + # Capture has to hit the bucket exactly or the graph is keyed on a total no + # runtime batch produces, and every ragged step silently runs eager. + tiers = [1, 3, 5] + for bs in (1, 2, 8): + for t in tiers: + bucket = bs * (t + 1) + got = _fit([1] * bs, bucket=bucket, padded_bs=bs, max_verify_len=1 + tiers[-1]) + assert int(got.verify_lens.sum()) == bucket + + +def _pad_len(*, bucket, n_real, n_pad, max_verify_len, floor=None): + """Mirror of the pad-row window ModelEngine.fit_ragged_verify_lens derives. + Returns None when no integer window leaves the real rows a reachable target.""" + real_capacity = n_real * max_verify_len + floor = n_real if floor is None else floor + if n_pad == 0: + pad_len = 1 + else: + lo = max(1, -(-(bucket - real_capacity) // n_pad)) # ceil division + hi = min(max_verify_len, (bucket - floor) // n_pad) + if lo > hi: + return None + pad_len = lo + real_target = bucket - n_pad * pad_len + if real_target < floor or real_target > real_capacity: + return None + return pad_len + + +@pytest.mark.parametrize( + "bucket, expected_pad_len", + [ + # 3 real rows capped at 6 hold up to 18; bucket 18 minus 5 one-token + # pad rows leaves 13, which the real rows absorb -- pads stay minimal. + (18, 1), + # The case that used to have no solution: bucket 24 with 5 one-token + # pad rows would demand 19 from real rows that cap at 18, so the pad + # rows must absorb what the real rows cannot. + (24, 2), + ], +) +def test_pad_rows_grow_when_the_real_rows_cannot_absorb_the_bucket(bucket, expected_pad_len): + padded_bs, max_verify_len = 8, 6 + real = [2, 3, 2] + n_pad = padded_bs - len(real) + pad_len = _pad_len(bucket=bucket, n_real=len(real), n_pad=n_pad, max_verify_len=max_verify_len) + assert pad_len == expected_pad_len + + real_target = bucket - n_pad * pad_len + assert sum(real) <= real_target <= len(real) * max_verify_len + got = _fit(real, bucket=real_target, padded_bs=len(real), max_verify_len=max_verify_len) + # The total the forward sees must be exactly the captured bucket. + assert int(got.verify_lens.sum()) + n_pad * pad_len == bucket + + +def test_pad_length_never_yields_an_infeasible_real_target(): + """Over every shape the scheduler can produce, the derived split must either + reconstruct the bucket exactly or decline outright (never a missed bucket).""" + max_verify_len = 6 + checked = declined = 0 + for padded_bs in (1, 2, 4, 8, 16, 32, 64): + grid = sorted({padded_bs * (t + 1) for t in range(max_verify_len)}) + for n_real in range(1, padded_bs + 1): + n_pad = padded_bs - n_real + for floor in range(n_real, n_real * max_verify_len + 1): + # The scheduler rounds up from "what the batch already needs + # plus one token per pad row". + fits = [b for b in grid if b >= floor + n_pad] + if not fits: + continue + bucket = fits[0] + checked += 1 + pad_len = _pad_len( + bucket=bucket, + n_real=n_real, + n_pad=n_pad, + max_verify_len=max_verify_len, + floor=floor, + ) + if pad_len is None: + declined += 1 + continue + real_target = bucket - n_pad * pad_len + assert floor <= real_target <= n_real * max_verify_len, ( + f"padded_bs={padded_bs} n_real={n_real} bucket={bucket} " + f"floor={floor} pad_len={pad_len} -> {real_target}" + ) + assert real_target + n_pad * pad_len == bucket + + assert checked > 10000, f"sweep too small to be meaningful: {checked}" + # Declining is safe (the step stays uniform) but should stay rare, or the + # bucket grid is too coarse to be worth capturing. + assert declined / checked < 0.15, f"{declined}/{checked} shapes have no valid split" + + +def test_extra_peer_stat_fields_are_ignored_here(): + """peer_stats may grow extra elements consumed by the caller; a rigid + two-element unpack here would raise on the first ragged step.""" + two = _shape(num_real=3, total=12, peers=[[3, 12], [4, 16]]) + three = _shape(num_real=3, total=12, peers=[[3, 12, 1], [4, 16, 1]]) + assert (three.padded_bs, three.bucket) == (two.padded_bs, two.bucket) + + +# --------------------------------------------------------------------------- +# device fill / row-map parity: fill_bucket_device and build_row_maps_device +# must reproduce the host path token-for-token. +# --------------------------------------------------------------------------- + + +def _host_fill(lens, padded_bs, bucket, max_verify_len): + layout = RaggedVerifyLayout.from_verify_lens( + torch.tensor(lens, dtype=torch.int32), + graph_num_tokens=bucket, + total_verify_tokens=sum(lens), + ) + return layout.fill_bucket(max_verify_len=max_verify_len, + padded_bs=padded_bs) + + +def _device_fill(lens, padded_bs, bucket, max_verify_len): + padded = torch.ones(padded_bs, dtype=torch.int32) + padded[:len(lens)] = torch.tensor(lens, dtype=torch.int32) + return fill_bucket_device( + padded, + num_real=torch.tensor(len(lens)), + graph_num_tokens=bucket, + max_verify_len=max_verify_len, + ) + + +def _feasible(lens, padded_bs, bucket, max_verify_len): + n_pad = padded_bs - len(lens) + return (sum(lens) + n_pad <= bucket <= padded_bs * max_verify_len) + + +def _all_lens(n_real, max_verify_len): + """Every [1, max]^n_real length vector.""" + if n_real == 0: + yield [] + return + for head in range(1, max_verify_len + 1): + for tail in _all_lens(n_real - 1, max_verify_len): + yield [head] + tail + + +@pytest.mark.parametrize("max_verify_len", [2, 4, 6]) +def test_fill_parity_exhaustive_small(max_verify_len): + """Exact host/device agreement over every feasible small case.""" + checked = 0 + for n_real in (1, 2, 3): + for padded_bs in (n_real, n_real + 1, n_real + 3): + for lens in _all_lens(n_real, max_verify_len): + for bucket in range(padded_bs, + padded_bs * max_verify_len + 1): + if not _feasible(lens, padded_bs, bucket, max_verify_len): + continue + host = _host_fill(lens, padded_bs, bucket, max_verify_len) + dev = _device_fill(lens, padded_bs, bucket, max_verify_len) + assert torch.equal(dev, host.verify_lens), ( + f"lens={lens} padded_bs={padded_bs} bucket={bucket} " + f"max={max_verify_len}: device {dev.tolist()} != " + f"host {host.verify_lens.tolist()}") + checked += 1 + assert checked > 100 + + +def test_fill_parity_randomized_large(): + """Production-shaped cases: bs up to 128, max_verify_len 6.""" + gen = torch.Generator().manual_seed(20260806) + max_verify_len = 6 + for _ in range(200): + n_real = int(torch.randint(1, 129, (1,), generator=gen)) + padded_bs = n_real + int(torch.randint(0, 9, (1,), generator=gen)) + lens = torch.randint(1, max_verify_len + 1, (n_real,), + generator=gen).tolist() + lo, hi = sum(lens) + (padded_bs - n_real), padded_bs * max_verify_len + bucket = int(torch.randint(lo, hi + 1, (1,), generator=gen)) + host = _host_fill(lens, padded_bs, bucket, max_verify_len) + dev = _device_fill(lens, padded_bs, bucket, max_verify_len) + assert torch.equal(dev, host.verify_lens) + assert int(dev.sum()) == bucket + + +def test_row_maps_match_prepare_semantics(): + """req_idx/correction must equal what prepare() stages: token o of a window + v gets correction o - v + 1, so the extent walks kv_len - v + 1 .. kv_len.""" + gen = torch.Generator().manual_seed(7) + for _ in range(50): + bs = int(torch.randint(1, 65, (1,), generator=gen)) + lens = torch.randint(1, 7, (bs,), generator=gen).to(torch.int32) + total = int(lens.sum()) + req_idx, corr = build_row_maps_device(lens, graph_num_tokens=total) + want_req, want_corr = [], [] + for r, v in enumerate(lens.tolist()): + for o in range(v): + want_req.append(r) + want_corr.append(o - v + 1) + assert req_idx.tolist() == want_req + assert corr.tolist() == want_corr + # Composing with a kv_lens gather is refresh_ragged_row_kv_lens: + # the extents must end exactly at each request's kv_len. + kv_lens = torch.randint(10, 1000, (bs,), generator=gen).to(torch.int32) + extents = kv_lens[req_idx] + corr + ends = torch.cumsum(lens.to(torch.long), 0) - 1 + assert torch.equal(extents[ends], kv_lens) + assert int(extents.min()) >= 1 + + +def test_fill_pad_fill_constrained_split(): + """With pad_fill given, every pad row carries exactly that many tokens, the + real rows absorb all slack, and the total still hits the bucket.""" + gen = torch.Generator().manual_seed(3) + max_verify_len = 6 + for _ in range(100): + n_real = int(torch.randint(1, 65, (1,), generator=gen)) + n_pad = int(torch.randint(0, 9, (1,), generator=gen)) + padded_bs = n_real + n_pad + pad_fill = int(torch.randint(1, max_verify_len + 1, (1,), + generator=gen)) + lens = torch.randint(2, max_verify_len + 1, (n_real,), generator=gen) + lo = int(lens.sum()) + n_pad * pad_fill + hi = n_real * max_verify_len + n_pad * pad_fill + if lo > hi: + continue + bucket = int(torch.randint(lo, hi + 1, (1,), generator=gen)) + padded = torch.ones(padded_bs, dtype=torch.int32) + padded[:n_real] = lens.to(torch.int32) + filled = fill_bucket_device(padded, + num_real=torch.tensor(n_real), + graph_num_tokens=bucket, + max_verify_len=max_verify_len, + pad_fill=pad_fill) + assert int(filled.sum()) == bucket + assert (filled[n_real:] == pad_fill).all() + assert (filled[:n_real] >= lens.to(torch.int32)).all() + assert int(filled[:n_real].max()) <= max_verify_len + + +def test_select_windows_respects_published_split(): + """With pad_len the prologue must land on the host fit's real/pad token + split exactly: real tokens = bucket - n_pad * pad_len, pads = pad_len.""" + gen = torch.Generator().manual_seed(41) + K = 5 + cfg = DSparkScheduleConfig(block_size=K, min_verify_len=1) + max_tok = cfg.resolved_max_verify_len + 1 + for _ in range(30): + n_real = int(torch.randint(1, 33, (1,), generator=gen)) + n_pad = int(torch.randint(0, 5, (1,), generator=gen)) + padded_bs = n_real + n_pad + pad_len = int(torch.randint(1, max_tok + 1, (1,), generator=gen)) + real_target = int( + torch.randint(n_real * 2, n_real * max_tok + 1, (1,), + generator=gen)) + bucket = real_target + n_pad * pad_len + budget = max(0, min(int(torch.randint(0, n_real * K + 1, (1,), + generator=gen)), + real_target - n_real * 2)) + conf = torch.randn(padded_bs + 4, K, generator=gen) * 2 + slot_idx = torch.arange(padded_bs, dtype=torch.long) + res = select_windows_device( + confidence_logits=conf, + slot_idx=slot_idx, + num_real=torch.tensor(n_real), + budget=torch.tensor(budget), + graph_num_tokens=bucket, + cfg=cfg, + pad_len=pad_len, + ) + assert int(res.verify_lens.sum()) == bucket + assert int(res.verify_lens[:n_real].sum()) == real_target + assert (res.verify_lens[n_real:] == pad_len).all() + assert int(res.verify_lens[:n_real].min()) >= cfg.min_verify_len + 1 + assert int(res.verify_lens[:n_real].max()) <= max_tok + + +def test_topk_tensor_budget_matches_int(): + """The 0-d tensor budget path (in-graph capacity knob) must allocate + identically to the host int path, including budget <= 0.""" + gen = torch.Generator().manual_seed(11) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + for _ in range(30): + bs = int(torch.randint(1, 65, (1,), generator=gen)) + survival = compute_survival( + torch.rand(bs, 5, generator=gen)) + for budget in (0, 1, bs, 2 * bs, bs * 4, bs * 10): + host = schedule_verify_lens_topk(survival=survival, + budget=budget, cfg=cfg) + dev = schedule_verify_lens_topk(survival=survival, + budget=torch.tensor(budget), + cfg=cfg) + assert torch.equal(host, dev) + + +def test_planner_lag2_ring_and_budget_split(monkeypatch): + """Device-window mode: the budget reads the snapshot one staging older + (lag-2), and the shape lens spread the budget uniformly to floor + budget.""" + monkeypatch.setenv("TLLM_DSPARK_DEVICE_WINDOWS", "1") + from tensorrt_llm._torch.speculative.dspark_verify import \ + DSparkVerifyPlanner + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + planner = DSparkVerifyPlanner(cfg=cfg) + assert planner.device_windows + + snap_a = torch.full((4, 5), 1.0) + snap_b = torch.full((4, 5), 2.0) + snap_c = torch.full((4, 5), 3.0) + planner.stage_confidence(snap_a) + assert planner._ready_older_snapshot() is None + planner.stage_confidence(snap_b) + older = planner._ready_older_snapshot() + assert older is not None and float(older[0, 0]) == 1.0 + planner.stage_confidence(snap_c) + older = planner._ready_older_snapshot() + assert float(older[0, 0]) == 2.0 + # The current snapshot is the freshest staging, untouched by the ring. + assert float(planner._ready_snapshot()[0, 0]) == 3.0 + + # Budget arithmetic through the frac override (snapshot-independent): + # shape lens sum to floor + budget and stay inside the bounds. + planner._forced_budget_frac = 0.5 + n = 7 + trimmable = cfg.resolved_max_verify_len - cfg.min_verify_len + budget, lens = planner.decide_verify_budget(num_gen_requests=n) + assert budget == int(round(0.5 * n * trimmable)) + assert len(lens) == n + assert sum(lens) == n * cfg.min_verify_len + budget + assert min(lens) >= cfg.min_verify_len + assert max(lens) <= cfg.resolved_max_verify_len + + +def _host_select(rows, budget, bucket, padded_bs, cfg): + """The lag-free host reference: calibrate -> survival -> topk -> fill.""" + survival = compute_survival(torch.sigmoid(rows)) + scheduled = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=cfg) + token_lens = (scheduled + 1).tolist() + layout = RaggedVerifyLayout.from_verify_lens( + torch.tensor(token_lens, dtype=torch.int32), + graph_num_tokens=bucket, + total_verify_tokens=sum(token_lens)) + return layout.fill_bucket( + max_verify_len=cfg.resolved_max_verify_len + 1, padded_bs=padded_bs) + + +def test_select_windows_device_matches_host_chain(): + """The full prologue equals the host chain run on the same (fresh) rows, + across pad rows, indirection through slots, and stale-stamp fallback.""" + gen = torch.Generator().manual_seed(20260806) + K = 5 + cfg = DSparkScheduleConfig(block_size=K, min_verify_len=1) + max_token_len = cfg.resolved_max_verify_len + 1 + for _ in range(40): + n_real = int(torch.randint(1, 33, (1,), generator=gen)) + padded_bs = n_real + int(torch.randint(0, 5, (1,), generator=gen)) + num_slots = padded_bs + 8 + conf = torch.randn(num_slots, K, generator=gen) * 3 + # Requests sit at shuffled slots; pad rows point at slot 0. + slots = torch.randperm(num_slots, generator=gen)[:n_real] + slot_idx = torch.zeros(padded_bs, dtype=torch.long) + slot_idx[:n_real] = slots + # A third of the rows are stale: their gather must neutralize to + # "verify everything". + stamp = torch.full((num_slots,), 7, dtype=torch.int64) + stale_rows = torch.rand(n_real, generator=gen) < 0.33 + stamp[slots[stale_rows]] = 3 + + budget = int(torch.randint(0, n_real * K, (1,), generator=gen)) + lo = n_real * (cfg.min_verify_len + 1) + min( + budget, n_real * (K - cfg.min_verify_len)) + (padded_bs - n_real) + bucket = min(padded_bs * max_token_len, + lo + int(torch.randint(0, 8, (1,), generator=gen))) + + result = select_windows_device( + confidence_logits=conf, + slot_idx=slot_idx, + num_real=torch.tensor(n_real), + budget=torch.tensor(budget), + graph_num_tokens=bucket, + cfg=cfg, + stamp=stamp, + expected_stamp=torch.tensor(7), + ) + rows = conf[slots].clone() + rows[stale_rows] = 30.0 # NEUTRAL_CONFIDENCE_LOGIT + host = _host_select(rows, budget, bucket, padded_bs, cfg) + assert torch.equal(result.verify_lens, host.verify_lens) + assert torch.equal(result.qo_indptr, host.qo_indptr) + assert int(result.verify_lens.sum()) == bucket + # Row maps must describe exactly the filled lens. + want_req, want_corr = build_row_maps_device( + host.verify_lens, graph_num_tokens=bucket) + assert torch.equal(result.req_idx, want_req) + assert torch.equal(result.kv_correction, want_corr) + + +def test_prologue_kv_pairing_matches_host_staging(): + """Rebasing the layout onto the true windows w needs kv_lens += 2*(w - S) + and offsets -= (w - S); any other split corrupts the indexer's slot mapping.""" + gen = torch.Generator().manual_seed(20260807) + for _ in range(20): + padded_bs = int(torch.randint(2, 40, (1,), generator=gen)) + past_seen = torch.randint(10, 500, (padded_bs,), generator=gen) + new_tokens_lens = torch.randint(1, 7, (padded_bs,), generator=gen) + shape_lens = torch.randint(2, 7, (padded_bs,), generator=gen) + true_lens = torch.randint(2, 7, (padded_bs,), generator=gen) + + # Host staging under the shape split S: the window counts twice. + kv_lens = (past_seen + 2 * shape_lens).to(torch.int32) + offsets = (new_tokens_lens - shape_lens).to(torch.int32) + + # Prologue: 2x on the kv side, 1x on the offsets side. + window_delta = (true_lens - shape_lens).to(torch.int32) + kv_lens += 2 * window_delta + offsets -= window_delta + + # Pre-offset kv_lens must equal host-with-w staging: the indexer's + # slot mapping and the expanded per-token extents derive from it. + assert torch.equal(kv_lens, + (past_seen + 2 * true_lens).to(torch.int32)) + + # In-graph correction at replay. + effective = kv_lens + offsets + + # Reference: host staging directly with the true windows w. + want = (past_seen + 2 * true_lens).to(torch.int32) + ( + new_tokens_lens - true_lens).to(torch.int32) + assert torch.equal(effective, want) + + +# -------------------------------------------------------------------------- +# overlap-scheduler input layout: the strided gather assumption and its +# ragged replacement, with the uniform batch reproducing the strided layout. +# -------------------------------------------------------------------------- + + +class _Req: + """Minimal stand-in exposing only what the layout code reads.""" + + def __init__(self, verify_len=None): + if verify_len is not None: + self.py_verify_len = verify_len + + +_MISSING = object() + + +@pytest.mark.parametrize( + "verify_len, runtime_tokens, expected", + [ + # Every non-DSpark request lands here: the attribute does not exist. + (_MISSING, 4, 4), + # The ragged path declined this request: the attribute is None. + (None, 4, 4), + # py_verify_len counts draft positions, matching runtime_draft_len; + # the accepted token makes it one more. + (2, 6, 3), + # A zero-draft window still carries the accepted token. + (0, 6, 1), + ], +) +def test_tokens_per_gen_step_resolves_each_request_window( + verify_len, runtime_tokens, expected): + req = _Req() + if verify_len is not _MISSING: + req.py_verify_len = verify_len + assert get_request_tokens_per_gen_step(req, runtime_tokens) == expected + + +def test_uniform_counts_reproduce_the_strided_gather(): + # This is the safety property the whole change rests on: with equal counts + # the index list picks exactly the elements tensor[slots, :width] does. + table = torch.arange(6 * 5).reshape(6, 5) + slots = [3, 0, 4] + width = 5 + + rows, cols = ragged_gather_index_lists(slots, [width] * len(slots)) + ragged = table[torch.tensor(rows), torch.tensor(cols)] + strided = table[torch.tensor(slots), :width].flatten() + + assert torch.equal(ragged, strided) + + +def test_ragged_counts_pack_each_request_window(): + table = torch.arange(4 * 5).reshape(4, 5) + rows, cols = ragged_gather_index_lists([2, 0], [3, 1]) + gathered = table[torch.tensor(rows), torch.tensor(cols)] + + # slot 2 columns 0..2, then slot 0 column 0 + assert gathered.tolist() == [10, 11, 12, 0] + + # A zero count contributes nothing while its neighbours keep their order, + # and the empty batch degenerates cleanly. + assert ragged_gather_index_lists([1, 2, 3], [2, 0, 1]) == ([1, 1, 3], [0, 1, 0]) + assert ragged_gather_index_lists([], []) == ([], []) + + +def test_mismatched_lengths_raise(): + with pytest.raises(ValueError, match="same length"): + ragged_gather_index_lists([0, 1], [3]) + with pytest.raises(ValueError, match="negative gather count"): + ragged_gather_index_lists([0], [-1]) + + +# -------------------------------------------------------------------------- +# compressor gen counts: _sync_gen_tokens_per_seq must source its per-request +# vector from the DEVICE seq-lens buffer (D2D, not a captured H2D of the host +# split) and use the global max window as the next_n scalar bound. +# -------------------------------------------------------------------------- + + +def _dummy(nc, staged_split, seq_lens_dev, max_draft_tokens=5, cap=16): + d = types.SimpleNamespace() + d.is_ragged_verify = True + d.num_generations = len(staged_split) + d.num_contexts = nc + d.max_draft_tokens = max_draft_tokens + # What the host shape split says (S); the fixed code must NOT read it. + d.ragged_verify_lens = list(staged_split) + d.gen_new_tokens_per_seq_cuda = torch.zeros(cap, dtype=torch.int) + d._seq_lens_cuda = torch.tensor(seq_lens_dev, dtype=torch.int) + return d + + +def _call(d, num_gen_tokens): + return DeepseekV4TrtllmAttentionMetadata._sync_gen_tokens_per_seq( + d, num_gen_tokens) + + +def test_vector_sources_from_seq_lens_not_host_split(): + # Simulate the prologue: host staged S, device rewrote seq lens to w. + S = [3, 3, 3] + w = [2, 4, 3] + d = _dummy(nc=0, staged_split=S, seq_lens_dev=w) + _call(d, sum(w)) + assert d.gen_new_tokens_per_seq.tolist() == w + + +def test_scalar_is_global_bound_not_batch_max(): + S = [2, 2] + d = _dummy(nc=0, staged_split=S, seq_lens_dev=S) + next_n = _call(d, sum(S)) + assert next_n == 1 + d.max_draft_tokens + assert d.num_gen_tokens_per_seq == 1 + d.max_draft_tokens + + +def test_gen_rows_slice_skips_context_rows(): + ctx_lens = [17, 9] + w = [4, 2] + d = _dummy(nc=2, staged_split=w, seq_lens_dev=ctx_lens + w) + _call(d, sum(w)) + assert d.gen_new_tokens_per_seq.tolist() == w + + +def test_uniform_branch_unchanged(): + d = _dummy(nc=0, staged_split=[3, 3], seq_lens_dev=[3, 3]) + d.is_ragged_verify = False + next_n = _call(d, 6) + assert next_n == 3 + assert d.gen_new_tokens_per_seq is None + + +# -------------------------------------------------------------------------- +# ragged CUDA-graph key: the key's token axis must describe the batch the +# graph will actually replay over (fit, publication, and the padding ladder). +# -------------------------------------------------------------------------- + + +class _Runner: + """Just the fields will_pad_to reads.""" + + def __init__(self, supported, cfg_batch_size): + self.supported_batch_sizes = supported + self.config = type("C", (), {"batch_size": cfg_batch_size})() + + +def test_will_pad_to_rejects_row_counts_the_ladder_cannot_reach(): + """The ragged fit must not derive a bucket grid from row counts the + captured padding ladder cannot reach.""" + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ + CUDAGraphRunner + + runner = _Runner([1, 2, 4, 8, 128, 192, 256], cfg_batch_size=256) + will = CUDAGraphRunner.will_pad_to + + assert will(runner, 256, 193), "193 -> 256 is reachable and should be taken" + assert not will(runner, 193, 193), "193 is not a captured size" + assert will(runner, 192, 192), "already at a captured size" + # Padding past the configured batch size is refused by _get_padded_batch, + # so the fit must not assume it. + assert not will(_Runner([256, 512], cfg_batch_size=256), 512, 300) + assert not will(runner, 0, 10) + assert not will(runner, 256, 0) + + +def test_graph_key_uses_the_agreed_bucket_not_a_fresh_sum(): + """Every attention-DP rank must key on the same fitted bucket -- never a + fresh re-sum of the padded batch, which can diverge across ranks.""" + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ + CUDAGraphRunner + + class _R: + spec_config = type("S", (), {"enable_ragged_verify": True})() + agreed_ragged_bucket = None + + def _batch(*verify_lens): + return types.SimpleNamespace(generation_requests=[ + types.SimpleNamespace(py_verify_len=v) for v in verify_lens + ]) + + r = _R() + get = CUDAGraphRunner._ragged_verify_bucket + windowed = _batch(3, 2, 1) + + # Nothing fitted yet -> not ragged, key unchanged from the uniform one. + assert get(r, windowed) is None + + # A fitted bucket is returned verbatim: the batch's own token total is + # never re-summed here, which is what makes every rank agree. + r.agreed_ragged_bucket = 449 + assert get(r, windowed) == 449 + # ... including when that total plainly differs from the fitted value. + assert get(r, _batch(1, 1)) == 449 + + # But a request without a window means the token-major staging chain never + # ran this step, so a ragged key would replay the PREVIOUS ragged step's + # row maps -- the ragged IMA. Decline, whatever the fit published. + assert get(r, _batch(3, None, 1)) is None + + # Cleared between steps so a stale bucket cannot key into the wrong graph. + r.agreed_ragged_bucket = None + assert get(r, windowed) is None + + # Config gate still wins: a non-ragged engine keeps the original key shape. + r.agreed_ragged_bucket = 512 + r.spec_config = type("S", (), {"enable_ragged_verify": False})() + assert get(r, windowed) is None + + +def test_capture_publishes_the_bucket_it_shaped_the_batch_to(): + """Whoever shapes a warmup batch must publish its bucket, or the graph is + keyed as if the batch were uniform and matches no runtime step.""" + from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + + runner = types.SimpleNamespace(agreed_ragged_bucket=None) + requests = [types.SimpleNamespace(py_verify_len=None) for _ in range(4)] + batch = types.SimpleNamespace(generation_requests=requests) + engine = types.SimpleNamespace(cuda_graph_runner=runner) + + # bs=4, draft_len=5 -> the t=5 bucket is 4 * 6 = 24 tokens. + PyTorchModelEngine._set_warmup_ragged_windows(engine, batch, 24, 5) + + assert sum(1 + r.py_verify_len for r in requests) == 24, ( + "the warmup batch must hit the bucket exactly") + assert runner.agreed_ragged_bucket == 24, ( + "capture shaped the batch to 24 tokens but did not publish it, so the " + "graph would be keyed as if the batch were uniform") + + # A narrower bucket at the same batch size, to prove it is per-entry and + # not set once: capture walks every tier for every batch size. + PyTorchModelEngine._set_warmup_ragged_windows(engine, batch, 8, 5) + assert sum(1 + r.py_verify_len for r in requests) == 8 + assert runner.agreed_ragged_bucket == 8 + + # An empty batch publishes nothing rather than a stale or zero bucket. + runner.agreed_ragged_bucket = None + PyTorchModelEngine._set_warmup_ragged_windows( + engine, types.SimpleNamespace(generation_requests=[]), 24, 5) + assert runner.agreed_ragged_bucket is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py new file mode 100644 index 000000000000..aa223abcb9b9 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py @@ -0,0 +1,668 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the DSpark SPS cost-table profiler (CPU only, no model load). + +Covers stats-row alignment, cell summarizing, the additive alpha/theta fit, +staircase compression, payload assembly/refusals, sweep geometry guards, the +profitability probe, and reading samples out of production iteration logs. +""" + +import importlib.util +import json +import pathlib +import sys +import tempfile + +import numpy as np +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable, total_verify_tokens + +# The profiler is a development tool under tests/microbenchmarks, not part of +# the shipped package; load it by path so this test follows the file. The +# module must be registered in sys.modules BEFORE exec so @dataclass can +# resolve annotations through sys.modules[cls.__module__].__dict__. +_SPEC = importlib.util.spec_from_file_location( + "dspark_sps_profiler", + pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / + "dspark_sps_profiler.py") +_PROFILER = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _PROFILER +_SPEC.loader.exec_module(_PROFILER) + +globals().update({name: getattr(_PROFILER, name) for name in ('CellStat', 'FlatCostTableError', 'InertCostTableError', 'InsufficientSamplesError', 'StepSample', 'SweepConfig', 'SweepGeometryError', 'aligned_steps_from_stats', 'build_cost_table_payload', 'check_table_is_informative', 'compress_to_risers', 'fit_additive_cost_model', 'load_cost_table', 'profitability_probe', 'resolve_padded_shape', 'running_max', 'samples_from_iter_log', 'summarize_cells')}) + +#: The exact set of keys ``dspark.py::_build_verify_planner`` reads. Anything +#: else in the file is ignored by the loader, so the profiler must not rely on +#: it -- and must not omit any of these. +LOADER_KEYS = { + "token_counts", + "step_time_ms", + "fixed_overhead_ms", + "batch_sizes", + "batch_overhead_ms", +} + + +def _stats_row(*, iteration, num_gen, rank=0, num_ctx=0, host_ms=10.0, gpu_ms=8.0): + """One /metrics-shaped iteration record.""" + row = { + "iter": iteration, + "attentionDpRank": rank, + "inflightBatchingStats": { + "numContextRequests": num_ctx, + "numGenRequests": num_gen, + }, + } + if host_ms is not None: + row["hostStepTimeMS"] = host_ms + if gpu_ms is not None: + row["gpuForwardTimeMS"] = gpu_ms + return row + + +# -------------------------------------------------------------------------- +# aligned step extraction +# -------------------------------------------------------------------------- + + +def test_aligned_steps_drops_prefill_iterations(): + """A step with context requests is a different shape, not a cheap decode.""" + steady = [_stats_row(iteration=i, num_gen=8, host_ms=10.0 + i) for i in range(3)] + assert aligned_steps_from_stats(steady) == [(0, 8, 10.0), (1, 8, 11.0), (2, 8, 12.0)] + rows = [ + _stats_row(iteration=0, num_gen=0, num_ctx=8, host_ms=90.0), + _stats_row(iteration=1, num_gen=8, num_ctx=2, host_ms=40.0), + _stats_row(iteration=2, num_gen=8, host_ms=10.0), + ] + assert aligned_steps_from_stats(rows) == [(2, 8, 10.0)] + + +def test_aligned_steps_requires_every_dp_rank(): + """Multi-rank alignment: iterations must be complete, agree on batch size, + and be timed by rank 0's clock, never averaged.""" + rows = [ + _stats_row(iteration=0, num_gen=8, rank=0), + _stats_row(iteration=0, num_gen=8, rank=1), + _stats_row(iteration=1, num_gen=8, rank=0), + ] + assert aligned_steps_from_stats(rows, expected_ranks=2) == [(0, 8, 10.0)] + + disagreeing = [ + _stats_row(iteration=0, num_gen=8, rank=0), + _stats_row(iteration=0, num_gen=7, rank=1), + ] + assert aligned_steps_from_stats(disagreeing, expected_ranks=2) == [] + + rank0_not_first = [ + _stats_row(iteration=0, num_gen=8, rank=1, host_ms=99.0), + _stats_row(iteration=0, num_gen=8, rank=0, host_ms=10.0), + ] + assert aligned_steps_from_stats(rank0_not_first, expected_ranks=2) == [(0, 8, 10.0)] + + +def test_aligned_steps_skips_missing_timing_field(): + """gpuForwardTimeMS is read without a sync and is absent on some steps.""" + rows = [ + _stats_row(iteration=0, num_gen=8, gpu_ms=None), + _stats_row(iteration=1, num_gen=8, gpu_ms=7.5), + ] + assert aligned_steps_from_stats(rows, timing_key="gpuForwardTimeMS") == [(1, 8, 7.5)] + + +# -------------------------------------------------------------------------- +# warmup + median +# -------------------------------------------------------------------------- + + +def _samples(batch_size, verify_len, times): + return [ + StepSample(batch_size=batch_size, verify_len=verify_len, step_time_ms=t, iteration=i) + for i, t in enumerate(times) + ] + + +def test_summarize_discards_warmup(): + warm = [500.0] * 4 + steady = [10.0] * 8 + cells = summarize_cells(_samples(8, 3, warm + steady), warmup_steps=4, min_samples=8) + assert len(cells) == 1 + assert cells[0].step_time_ms == pytest.approx(10.0) + assert cells[0].num_samples == 8 + + +def test_summarize_reports_median_not_mean(): + """The right tail is real and must not move the point.""" + times = [10.0, 10.0, 10.0, 10.0, 10.0, 1000.0] + cells = summarize_cells(_samples(8, 3, times), warmup_steps=0, min_samples=1) + assert cells[0].step_time_ms == pytest.approx(10.0) + assert cells[0].step_time_ms != pytest.approx(float(np.mean(times))) + # ...but the tail is still reported, so a pathological cell is visible. + assert cells[0].p90_ms > cells[0].step_time_ms + + +def test_summarize_drops_thin_cells_outside_the_requested_grid(): + """Thin incidental cells are dropped, thick ones kept, and a thin + REQUESTED cell must still refuse -- with or without a requested grid.""" + grid = [(8, 3)] + requested = _samples(8, 3, [10.0] * 8) + ramp = _samples(2, 3, [30.0]) # thin, incidental: drop + thick_extra = _samples(4, 3, [20.0] * 8) # thick, incidental: keep + cells = summarize_cells(requested + ramp + thick_extra, warmup_steps=0, + min_samples=4, expected_cells=grid) + assert [(c.batch_size, c.verify_len) for c in cells] == [(4, 3), (8, 3)] + with pytest.raises(InsufficientSamplesError, match="steady samples"): + summarize_cells(_samples(8, 3, [10.0]) + thick_extra, warmup_steps=0, + min_samples=4, expected_cells=grid) + # No grid given means every observed cell was requested: same refusal. + with pytest.raises(InsufficientSamplesError, match="steady samples"): + summarize_cells(_samples(8, 3, [10.0] * 5), warmup_steps=4, min_samples=8) + + +def test_summarize_orders_by_iteration_before_cutting_warmup(): + """Stats arrive per queue, not necessarily in iteration order.""" + out_of_order = [ + StepSample(batch_size=8, verify_len=1, step_time_ms=t, iteration=i) + for i, t in sorted(enumerate([500.0, 500.0, 10.0, 10.0]), key=lambda p: -p[0]) + ] + cells = summarize_cells(out_of_order, warmup_steps=2, min_samples=2) + assert cells[0].step_time_ms == pytest.approx(10.0) + + +def test_step_sample_json_round_trip(): + sample = StepSample(batch_size=16, verify_len=3, step_time_ms=12.5, iteration=7) + assert StepSample.from_json(json.loads(json.dumps(sample.to_json()))) == sample + assert sample.total_verify_tokens == total_verify_tokens(16, 3) == 64 + + +# -------------------------------------------------------------------------- +# additive fit +# -------------------------------------------------------------------------- + + +def _synthetic_cells(batch_sizes, verify_lens, *, alpha, theta, noise=None): + """Cells generated by an exactly additive machine ``alpha(bs) + theta(M)``.""" + cells = [] + for batch_size in batch_sizes: + for verify_len in verify_lens: + tokens = total_verify_tokens(batch_size, verify_len) + value = alpha(batch_size) + theta(tokens) + if noise is not None: + value += noise(batch_size, verify_len) + cells.append( + CellStat( + batch_size=batch_size, + verify_len=verify_len, + step_time_ms=value, + num_samples=32, + p10_ms=value, + p90_ms=value, + ) + ) + return cells + + +def _staircase(edges_and_costs): + def theta(tokens): + value = edges_and_costs[0][1] + for edge, cost in edges_and_costs: + if tokens >= edge: + value = cost + return value + + return theta + + +BATCH_SIZES = [8, 16, 32, 64] +VERIFY_LENS = [1, 2, 3, 4, 5] +ALPHA = {8: 20.0, 16: 22.0, 32: 26.0, 64: 34.0} +THETA = _staircase([(0, 1.0), (64, 2.0), (128, 4.0), (256, 8.0)]) + + +def test_fit_recovers_relative_structure_exactly(): + """The fit is only identified up to one shared constant -- pin that, then + every cell must be reproduced to the last bit.""" + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) + fit = fit_additive_cost_model(cells) + assert fit.warnings == () + + for cell in cells: + predicted = fit.intercept_ms[cell.batch_size] + fit.theta_ms[cell.total_verify_tokens] + assert predicted == pytest.approx(cell.step_time_ms, abs=1e-6) + assert fit.max_rel_residual == pytest.approx(0.0, abs=1e-9) + + # The batch axis keeps its measured *differences*, which is what alpha means. + for smaller, larger in zip(BATCH_SIZES, BATCH_SIZES[1:]): + measured = ALPHA[larger] - ALPHA[smaller] + fitted = fit.intercept_ms[larger] - fit.intercept_ms[smaller] + assert fitted == pytest.approx(measured, abs=1e-6) + + +def test_fit_surfaces_a_bad_cell_rather_than_hiding_it(): + """A cell the additive model cannot explain is surfaced in the warnings, + not dropped as noise.""" + + def noise(batch_size, verify_len): + return 500.0 if (batch_size, verify_len) == (32, 3) else 0.0 + + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA, noise=noise) + fit = fit_additive_cost_model(cells) + + assert fit.max_abs_residual_ms > 100.0 + assert fit.max_rel_residual > 1.0 + report = " ".join(fit.warnings) + assert "additive model misses" in report + # The least-squares solve spreads the bad cell over everything sharing its + # batch size or token count, so the report names a neighbourhood -- the bad + # cell is in it, but so are innocent ones, which is why it is worded that way. + assert "bs=32, L=3" in report + + +def test_fit_clamps_a_dip_upward(): + """A non-monotone theta would offer the planner a cheaper *longer* verify.""" + dipped = _staircase([(0, 1.0), (64, 5.0), (128, 3.0), (256, 8.0)]) + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=dipped) + fit = fit_additive_cost_model(cells) + + tokens = sorted(fit.theta_ms) + values = [fit.theta_ms[m] for m in tokens] + assert values == running_max(values) + assert any("monoton" in w for w in fit.warnings) + + +def test_fit_refuses_a_disconnected_sweep(): + """Two islands each carry their own unknown offset; splicing them lies.""" + cells = _synthetic_cells([8], [1, 2], alpha=ALPHA.get, theta=THETA) + cells += _synthetic_cells([100], [1, 2], alpha=lambda _: 40.0, theta=THETA) + with pytest.raises(SweepGeometryError, match="disconnected"): + fit_additive_cost_model(cells) + + +# -------------------------------------------------------------------------- +# staircase shaping +# -------------------------------------------------------------------------- + + +def test_running_max_is_monotone(): + assert running_max([1.0, 3.0, 2.0, 2.5, 9.0]) == [1.0, 3.0, 3.0, 3.0, 9.0] + + +def test_compress_keeps_risers_and_closes_their_shelves(): + """A shelf survives as a breakpoint PAIR so the interpolating consumer + stays flat where the hardware was flat.""" + tokens = [16, 32, 48, 64, 80] + times = [1.0, 1.002, 1.004, 2.0, 2.001] + kept_tokens, kept_times = compress_to_risers( + tokens, times, min_riser_ms=0.02, max_breakpoints=8 + ) + assert kept_tokens == [16, 48, 64] + assert kept_times == [1.0, 1.004, 2.0] + + +def test_truncation_drops_the_point_cheapest_to_interpolate(): + """Over budget, remove the interior point interpolation best reconstructs + -- never the shelf-side point, which would re-bill mid-shelf totals.""" + tokens = [16, 32, 64, 128] + times = [1.0, 1.1, 4.0, 9.0] + kept_tokens, _ = compress_to_risers(tokens, times, min_riser_ms=0.0, max_breakpoints=3) + assert kept_tokens == [16, 32, 128] + + +def test_compress_never_drops_the_base_shelf(): + """Index 0 is what every below-range query clamps onto.""" + tokens = [16, 32, 64] + times = [1.0, 100.0, 101.0] + kept_tokens, _ = compress_to_risers(tokens, times, min_riser_ms=0.0, max_breakpoints=2) + assert kept_tokens[0] == 16 + + +# -------------------------------------------------------------------------- +# payload assembly and JSON round-trip through SpsCostTable +# -------------------------------------------------------------------------- + + +def _payload(): + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) + return build_cost_table_payload(cells, riser_tolerance=0.01, max_breakpoints=8), cells + + +def test_round_trip_reproduces_measured_cells(): + """The payload carries exactly the loader's keys, survives a JSON round + trip into a real SpsCostTable, and prices every cell within one shelf.""" + payload, cells = _payload() + assert LOADER_KEYS.issubset(payload) + assert set(payload) - LOADER_KEYS == {"_meta"} + + table = load_cost_table(json.loads(json.dumps(payload))) + assert isinstance(table, SpsCostTable) + assert not table.is_flat + # The loader builds SpsCostTable directly, so the profiler's own validation + # must be the real constructor rather than a copy of its rules. + assert list(table.token_counts) == payload["token_counts"] + assert len(table.batch_sizes) == len(table.batch_overhead_ms) + + for cell in cells: + predicted = table.step_time(cell.total_verify_tokens, cell.batch_size) + assert predicted == pytest.approx(cell.step_time_ms, abs=1e-4) + + +def test_emitted_terms_are_not_double_counted(): + """step_time_ms is theta only; the floor lives in the two overhead fields.""" + payload, cells = _payload() + table = load_cost_table(payload) + smallest = min(BATCH_SIZES) + floor = payload["fixed_overhead_ms"] + table.batch_overhead(smallest) + # The floor is the M-independent part, i.e. the intercept of the measured + # row at the smallest batch size -- not the whole step time. + row = sorted((c for c in cells if c.batch_size == smallest), key=lambda c: c.verify_len) + assert floor < row[0].step_time_ms + assert floor == pytest.approx(ALPHA[smallest], abs=0.5) + + +def test_batch_overhead_grows_with_batch_size(): + payload, cells = _payload() + overheads = payload["batch_overhead_ms"] + assert overheads == sorted(overheads) + assert payload["batch_sizes"] == sorted(BATCH_SIZES) + assert payload["fixed_overhead_ms"] >= 0.0 + # Provenance rides along under _meta: the encoding marker and the cells. + assert payload["_meta"]["encoding"] == "decomposed" + assert len(payload["_meta"]["cells"]) == len(cells) + + +# -------------------------------------------------------------------------- +# refusals +# -------------------------------------------------------------------------- + + +def test_refuses_a_flat_curve(): + """A machine whose step cost ignores M must not produce a table at all.""" + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=lambda _: 3.0) + payload = build_cost_table_payload(cells) + assert load_cost_table(payload).is_flat + with pytest.raises(FlatCostTableError, match="is_flat"): + check_table_is_informative(payload, batch_sizes=BATCH_SIZES, tiers=[1, 3, 5]) + + +def test_refuses_a_curve_that_only_moves_by_noise(): + """is_flat is exact float equality, so it cannot catch this one.""" + cells = _synthetic_cells( + BATCH_SIZES, + VERIFY_LENS, + alpha=ALPHA.get, + theta=lambda tokens: 3.0 + 1e-4 * tokens, + ) + payload = build_cost_table_payload(cells, riser_tolerance=0.0, max_breakpoints=8) + assert not load_cost_table(payload).is_flat + with pytest.raises(FlatCostTableError, match="noise floor"): + check_table_is_informative( + payload, batch_sizes=BATCH_SIZES, tiers=[1, 3, 5], min_step_time_spread=0.02 + ) + + +# A real but shallow riser: at bs=32 tier 3 (M=128) sits below it and tier 5 +# (M=192) above, so the curve is genuinely non-flat and clears the noise floor +# (2.5%), yet the step only gets 2.5% more expensive while the acceptance yield +# given up between those tiers is ~36% at p=0.9. The argmax therefore never +# leaves max_tier -- non-flat and inert at the same time. The shelf is encoded +# as a breakpoint pair so the interpolating consumer keeps it flat. +INERT_PAYLOAD = { + "token_counts": [0, 159, 160], + "step_time_ms": [2.0, 2.0, 2.5], + "fixed_overhead_ms": 18.0, + "batch_sizes": [32], + "batch_overhead_ms": [0.0], +} + + +def test_refuses_an_inert_table(): + """Non-flat, but no tier pair is worth trading yield for; allow_inert=True + is the escape hatch and must still say that nothing trims.""" + kwargs = dict(batch_sizes=[32], tiers=[1, 3, 5], acceptance_rates=[0.9, 0.95]) + with pytest.raises(InertCostTableError, match="inert"): + check_table_is_informative(INERT_PAYLOAD, **kwargs) + diagnostics = check_table_is_informative(INERT_PAYLOAD, allow_inert=True, **kwargs) + assert diagnostics["trims_somewhere"] is False + + +def test_informative_table_passes_and_reports_where_trimming_wins(): + payload, _ = _payload() + diagnostics = check_table_is_informative( + payload, + batch_sizes=BATCH_SIZES, + tiers=[1, 3, 5], + acceptance_rates=[0.3, 0.6, 0.9], + ) + assert diagnostics["trims_somewhere"] is True + assert diagnostics["best_step_time_spread"] > 0.02 + # Even a good table cannot rescue a one-rung ladder: nothing to choose. + with pytest.raises(InertCostTableError, match="single entry"): + check_table_is_informative(payload, batch_sizes=BATCH_SIZES, tiers=[5]) + + +# -------------------------------------------------------------------------- +# sweep geometry guards (pure; the sweep itself needs a GPU) +# -------------------------------------------------------------------------- + + +def _sweep(**overrides) -> SweepConfig: + kwargs = dict( + model="/models/x", + speculative_model="/models/x", + batch_sizes=[8, 16], + verify_lens=[1, 3, 5], + input_len=1024, + warmup_steps=8, + measure_steps=32, + max_seq_len=4096, + max_num_tokens=8192, + ) + kwargs.update(overrides) + return SweepConfig(**kwargs) + + +def test_token_budget_is_the_step_count_when_acceptance_is_pinned(): + """Pinned, a step commits exactly one token, so the two are the same + number; unpinned it can commit verify_len + 1, so the budget must scale.""" + pinned = _sweep() + assert pinned.max_tokens_for(5) == pinned.warmup_steps + pinned.measure_steps + 16 + loose = _sweep(pin_acceptance=0.0) + assert loose.max_tokens_for(5) == pinned.max_tokens_for(5) * 6 + + +def test_validate_rejects_a_sequence_that_would_be_evicted(): + """An evicted request drains the batch mid-cell and shrinks the shape.""" + with pytest.raises(SweepGeometryError, match="max-seq-len"): + _sweep(input_len=4064, block_follows_verify_len=True).validate() + + +def test_validate_rejects_a_decode_step_wider_than_max_num_tokens(): + """A split decode step never has the shape its cell is filed under.""" + with pytest.raises(SweepGeometryError, match="widest decode step"): + _sweep(batch_sizes=[8, 4096], max_num_tokens=1024, max_seq_len=8192, + block_follows_verify_len=True).validate() + + +def test_validate_accepts_a_sane_sweep(): + """Sane includes verifying below a constant drafted block.""" + _sweep(block_follows_verify_len=True).validate() + _sweep(verify_lens=[5], max_draft_len=5).validate() + _sweep().validate() + + +def test_dp_size_only_counts_ranks_under_attention_dp(): + """Non-ADP publishes one stats row per iteration however wide TP is.""" + assert _sweep(tp_size=8).dp_size == 1 + assert _sweep(tp_size=8, enable_attention_dp=True).dp_size == 8 + + +# -------------------------------------------------------------------------- +# profitability probe (runs the real planner argmax) +# -------------------------------------------------------------------------- + + +def test_probe_trims_at_low_acceptance_and_not_at_high(): + """Trimming is a trade: it only wins when the drafts are unlikely to survive.""" + payload = { + # A riser between bs*(3+1)=128 and bs*(5+1)=192 at bs=32, so tiers 3 and + # 5 sit on different shelves and the argmax has something to choose: + # T(3) = 3 ms, T(5) = 4 ms, a 33% step for a yield gain that is worth it + # only when the drafts survive. + "token_counts": [0, 160], + "step_time_ms": [2.0, 3.0], + "fixed_overhead_ms": 0.0, + "batch_sizes": [32], + "batch_overhead_ms": [1.0], + } + probes = profitability_probe( + payload, batch_size=32, tiers=[1, 3, 5], acceptance_rates=[0.05, 0.99] + ) + by_rate = {p["acceptance_rate"]: p for p in probes} + assert by_rate[0.05]["chosen_verify_len"] < 5 + assert by_rate[0.05]["trims"] is True + assert by_rate[0.99]["chosen_verify_len"] == 5 + + +def test_probe_uses_the_bonus_token_convention(): + """A tier's cost must be looked up at bs*(L+1), never bs*L -- the same + convention as ``total_verify_tokens``.""" + payload = { + # Priced so that tier 5 (M = 8*6 = 48) is expensive and tier 3 + # (M = 8*4 = 32) is cheap. Under the wrong bs*L convention tier 5 would + # be M = 40 and land on the cheap shelf instead. The shelf is a + # breakpoint pair so the interpolating consumer keeps it flat. + "token_counts": [0, 43, 44], + "step_time_ms": [1.0, 1.0, 50.0], + "fixed_overhead_ms": 0.0, + "batch_sizes": [8], + "batch_overhead_ms": [0.0], + } + table = load_cost_table(payload) + assert table.step_time(total_verify_tokens(8, 5), 8) == pytest.approx(50.0) + assert table.step_time(total_verify_tokens(8, 3), 8) == pytest.approx(1.0) + probes = profitability_probe(payload, batch_size=8, tiers=[1, 3, 5], acceptance_rates=[0.9]) + assert probes[0]["chosen_verify_len"] == 3 + + +# -------------------------------------------------------------------------- +# reading cost-table samples out of production iteration logs +# -------------------------------------------------------------------------- + +# A deployment's captured ladder, abbreviated. Padding rounds a step's rows up +# to the smallest entry at or above them, which is what turns a padded token +# total into exactly one verify length. +LADDER = [1, 2, 4, 8, 16, 32, 64, 128, 192, 256] + + +def _line(iteration, bs, step_ms, gen_tokens, rank="0"): + return (f"{rank}: [08/04/2026-19:08:13] [TRT-LLM] [I] [_torch][RANK 0] " + f"iter = {iteration}, global_rank = 0, rank = 0, " + f"num_scheduled_requests = {bs}, kv_cache_util = 0.125, " + f"currank_total_requests = 1/8, host_step_time = {step_ms}ms, " + f"prev_device_step_time = {step_ms}ms, " + f"timestamp = 2026-08-04 19:08:13, " + f"states = {{'num_ctx_requests': 0, 'num_ctx_tokens': 0, " + f"'num_generation_tokens': {gen_tokens}, " + f"'cached_kv_tokens': 63821}}") + + +def _write(lines, name="gen.log"): + # tempfile rather than pytest's tmp_path: these tests also run under the + # fixture-less in-container runner used for the GPU images. + path = pathlib.Path(tempfile.mkdtemp()) / name + path.write_text("\n".join(lines) + "\n") + return str(path) + + +def test_verify_length_comes_from_the_padded_token_total(): + """The label is read out of the log, never inferred from step times; steps + the ladder cannot explain are skipped, not mislabelled.""" + log = _write([ + _line(1, 1, 62098.0, 6), # executor's first step: dropped + _line(2, 244, 79.2, 768), # 244 rows padded to 256 -> rung-2 + _line(3, 242, 117.1, 1536), # padded to 256 -> full block + _line(4, 128, 61.0, 384), # exactly 128 rows -> rung-2 + _line(5, 256, 300.0, 2048), # L = 7, past the block: skipped + _line(6, 300, 400.0, 1800), # more rows than the ladder holds: skipped + ]) + samples = samples_from_iter_log([log], max_draft_len=5, + padded_batch_sizes=LADDER) + got = sorted((s.batch_size, s.verify_len, s.step_time_ms) + for s in samples) + assert got == [(128, 2, 61.0), (256, 2, 79.2), (256, 5, 117.1)] + # The table is indexed by the padded token total, which each sample exposes. + assert sorted(s.total_verify_tokens for s in samples) == [384, 768, 1536] + + +def test_the_padded_width_is_what_gets_recorded(): + """A padded step is filed under the padded width, the ladder disambiguates + the verify length, and an unexplainable token total drops the step.""" + assert resolve_padded_shape(num_rows=244, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) == (256, 2) + assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256) is None + assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) == (128, 5) + assert resolve_padded_shape(num_rows=200, num_generation_tokens=600, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) is None + + +def test_the_first_step_of_each_executor_is_dropped(): + """Each executor numbers from iter 1 and its first step's host time + contains the wait for the first request, so it is never a measurement.""" + log = _write([ + _line(1, 1, 200.0, 6), # KV-estimation executor + _line(2, 256, 79.0, 768), + _line(1, 1, 62098.0, 6), # real executor starts over + _line(2, 256, 79.4, 768), + ]) + times = [s.step_time_ms for s in samples_from_iter_log([log], max_draft_len=5)] + assert times == [79.0, 79.4] + + +def test_only_the_requested_rank_is_read(): + """Every rank logs; counting all of them would multiply every cell. + Outlier steps are bounded out on the same pass.""" + log = _write([ + _line(2, 256, 79.2, 768, rank="0"), + _line(2, 256, 79.9, 768, rank="3"), # another rank's copy: dropped + _line(3, 256, 5000.0, 768), # outlier past max_step_ms + ]) + times = [s.step_time_ms for s in + samples_from_iter_log([log], max_draft_len=5, max_step_ms=1000.0)] + assert times == [79.2] + + +def test_several_logs_are_pooled(): + """One arm gives one rung; the table needs the ladder, so pool the runs.""" + a = _write([_line(1, 1, 100.0, 6), _line(2, 256, 117.1, 1536)], + name="a.log") + c = _write([_line(1, 1, 100.0, 6), _line(2, 256, 79.2, 768)], + name="c.log") + got = sorted(s.verify_len for s in + samples_from_iter_log([a, c], max_draft_len=5)) + assert got == [2, 5] + + +def test_a_log_without_the_states_dict_yields_nothing(): + """enable_iter_perf_stats off: fail loudly upstream, not with a bad table.""" + path = pathlib.Path(tempfile.mkdtemp()) / "bare.log" + path.write_text( + "0: [TRT-LLM] iter = 2, num_scheduled_requests = 256, " + "host_step_time = 79.2ms\n") + assert samples_from_iter_log([str(path)], max_draft_len=5) == [] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py new file mode 100644 index 000000000000..f9644d183e3d --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py @@ -0,0 +1,237 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark STS calibration subsystem: recorder ring pairing, shard provenance, +fitter acceptance, collection-regime guards, and confidence-head resolution. +Runs on CPU only (the recorder is device-agnostic), so it fits pre-merge CI. +""" + +import contextlib +import importlib.util +import os +import pathlib +import sys +import tempfile + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_sts import (DSparkStsRecorder, + STS_COLLECT_ENV, + make_recorder_from_env) + +BLOCK = 5 +ROWS = 8 + +_PIN = "TLLM_DSPARK_FORCE_VERIFY_LEN" + + +def _recorder(tmp, flush_every=1000): + return DSparkStsRecorder(path_stem=os.path.join(tmp, "shard"), + block_size=BLOCK, rank=0, + flush_every=flush_every) + + +def _stage(rec, seq, written_rows, fill): + """Stage one pass's snapshot: `written_rows` stamped `seq`, rest stale.""" + logits = torch.full((ROWS, BLOCK), float(fill)) + stamps = torch.zeros(ROWS, dtype=torch.int32) + stamps[list(written_rows)] = seq + rec.stage_snapshot(device_logits=logits, device_stamps=stamps, + staged_seq=seq) + + +def test_pairs_by_pass_identity_not_arrival_order(): + """A label for pass 1 gets pass-1 logits even after pass 2 was staged.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2], fill=1.0) + _stage(rec, seq=2, written_rows=[2], fill=2.0) # overwrite arrives first + rec.record(row=2, accepted=3, target_seq=1) + rec.record(row=2, accepted=1, target_seq=2) + assert rec.stats["recorded"] == 2 + assert rec._logits[0][0].item() == 1.0 # pass 1 content, not latest + assert rec._logits[1][0].item() == 2.0 + + +def test_row_not_written_for_that_pass_is_declined(): + """A row the pass never drafted must not contribute a fabricated pair.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2], fill=1.0) + rec.record(row=3, accepted=2, target_seq=1) # row 3 stamp is 0, not 1 + assert rec.stats["recorded"] == 0 + assert rec.stats["stale_stamp"] == 1 + + +def test_ring_eviction_declines_instead_of_mispairing(): + """A label older than the ring depth is dropped, never matched to the + evicting pass's content.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + for seq in range(1, 2 + DSparkStsRecorder._RING_DEPTH): + _stage(rec, seq=seq, written_rows=[2], fill=float(seq)) + rec.record(row=2, accepted=3, target_seq=1) # slot now holds seq 5 + assert rec.stats["recorded"] == 0 + assert rec.stats["stale_stamp"] == 1 + + +def test_every_giving_up_path_is_counted(): + """No decline path may be silent; every one increments a named counter.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + rec.record(row=2, accepted=1, target_seq=None) # no pass id + rec.record(row=2, accepted=1, target_seq=7) # ring slot never staged + _stage(rec, seq=1, written_rows=[2], fill=1.0) + rec.record(row=None, accepted=1, target_seq=1) # unresolvable request + rec.record(row=ROWS + 3, accepted=1, target_seq=1) + assert rec.stats == {"recorded": 0, "no_snapshot": 2, + "stale_stamp": 0, "row_out_of_range": 1, + "no_row": 1, "snapshots_staged": 1} + + +def test_staging_without_stamps_is_a_noop(): + """No stamps means no way to verify identity -- refuse to guess.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + rec.stage_snapshot(device_logits=torch.zeros(ROWS, BLOCK), + device_stamps=None, staged_seq=1) + rec.stage_snapshot(device_logits=torch.zeros(ROWS, BLOCK), + device_stamps=torch.zeros(ROWS, dtype=torch.int32), + staged_seq=None) + assert rec.stats["snapshots_staged"] == 0 + rec.record(row=2, accepted=1, target_seq=1) + assert rec.stats["no_snapshot"] == 1 + + +def test_shard_carries_provenance_and_prefix_semantics(): + """Flushed shards carry ring provenance, stats, and prefix-mask labels.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2], fill=1.5) + rec.record(row=2, accepted=3, target_seq=1) + rec.flush() + blob = torch.load(os.path.join(tmp, "shard.r0.0.pt"), + map_location="cpu") + assert blob["meta"]["pairing"] == "draft_seq_ring" + assert blob["meta"]["stats"]["recorded"] == 1 + assert blob["logits"].shape == (1, BLOCK) + # accepted=3 -> positions 0..2 were ALL accepted: the prefix event + # whose probability survival[k] estimates. + assert blob["prefix_mask"].tolist() == [[1.0, 1.0, 1.0, 0.0, 0.0]] + + +_SPEC = importlib.util.spec_from_file_location( + "dspark_fit_sts", + pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / + "dspark_fit_sts.py") +_FITTER = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _FITTER +_SPEC.loader.exec_module(_FITTER) + + +def test_fitter_refuses_shards_without_ring_provenance(): + """The fitter must refuse shards lacking draft_seq_ring provenance.""" + with tempfile.TemporaryDirectory() as tmp: + torch.save({"logits": torch.zeros(4, BLOCK), + "prefix_mask": torch.zeros(4, BLOCK)}, + os.path.join(tmp, "old.r0.0.pt")) + with pytest.raises(SystemExit, match="draft_seq_ring"): + _FITTER.load_shards(os.path.join(tmp, "*.pt")) + + +def test_fitter_accepts_ring_shards(): + """Shards produced by the ring recorder load cleanly in the fitter.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2, 3], fill=0.5) + rec.record(row=2, accepted=2, target_seq=1) + rec.record(row=3, accepted=5, target_seq=1) + rec.flush() + logits, mask = _FITTER.load_shards(os.path.join(tmp, "*.pt")) + assert logits.shape == (2, BLOCK) and mask.shape == (2, BLOCK) + + +@contextlib.contextmanager +def collecting(pin=None): + """Set the collection env for the block, restore it after.""" + saved = {k: os.environ.get(k) for k in (STS_COLLECT_ENV, _PIN)} + try: + with tempfile.TemporaryDirectory() as tmp: + os.environ[STS_COLLECT_ENV] = os.path.join(tmp, "collect") + if pin is None: + os.environ.pop(_PIN, None) + else: + os.environ[_PIN] = str(pin) + yield + finally: + for k, v in saved.items(): + os.environ.pop(k, None) if v is None else os.environ.__setitem__(k, v) + + +def test_collection_off_builds_nothing(): + """The serving default must be untouched by any of this.""" + os.environ.pop(STS_COLLECT_ENV, None) + assert make_recorder_from_env(block_size=5) is None + assert make_recorder_from_env(block_size=5, has_cost_table=True, + ragged_mode="compact") is None + + +def test_a_loaded_cost_table_without_a_pin_is_refused(): + """A table means the planner trims, and trimming censors the label.""" + with collecting(): + with pytest.raises(ValueError, match="censors the acceptance label"): + make_recorder_from_env(block_size=5, has_cost_table=True) + + +def test_a_pinned_window_makes_the_table_harmless(): + """Every uncensored regime (pinned window, no table, static mode) must + build a recorder.""" + with collecting(pin=5): + rec = make_recorder_from_env(block_size=5, has_cost_table=True) + assert rec is not None and rec.block_size == 5 + with collecting(): + assert make_recorder_from_env(block_size=5, + has_cost_table=False) is not None + assert make_recorder_from_env(block_size=5, + ragged_mode="static") is not None + + +def test_compact_mode_is_refused(): + """Padding rows contribute prefix labels that measure nothing; refused + even with the pin set.""" + with collecting(pin=5): + with pytest.raises(ValueError, match="Padded verify rows"): + make_recorder_from_env(block_size=5, ragged_mode="compact") + + +class _Head: + pass + + +class _Bare: + """DSparkDraftModel layout: stages under .mtp_layers.""" + + def __init__(self, head): + stage = type("Stage", (), {})() + if head is not None: + stage.confidence_head = head + self.mtp_layers = [object(), object(), stage] + + +class _Wrapper: + """DSparkForCausalLM layout: bare model under .dspark_model.""" + + def __init__(self, head): + self.dspark_model = _Bare(head) + + +def test_head_resolves_across_both_layouts_and_absence_is_none(): + """resolve_confidence_head finds the head under either model layout and + returns None when absent.""" + from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head + h = _Head() + assert resolve_confidence_head(_Bare(h)) is h + assert resolve_confidence_head(_Wrapper(h)) is h + assert resolve_confidence_head(_Bare(None)) is None + assert resolve_confidence_head(_Wrapper(None)) is None + assert resolve_confidence_head(object()) is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 4625d1b0115a..ba186f4d942e 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -58,46 +58,40 @@ def test_metadata_buffer_and_layer_lookup(): # sorted, O(1) membership assert meta.is_layer_capture(58) and meta.is_layer_capture(60) assert not meta.is_layer_capture(0) and not meta.is_layer_capture(61) + # prepare() publishes contiguous batch indices for the current batch. + meta.request_ids = [7, 3, 5] + meta.prepare() + assert meta.batch_indices_cuda[:3].tolist() == [0, 1, 2] def test_metadata_capture_plain_hidden(): - """A [num_tokens, hidden] capture is stored at the layer's slice as-is.""" + """Capture routing: a [num_tokens, hidden] capture is stored at the layer's + slice as-is; a flattened mHC residual [N, hc_mult*hidden] is reduced by mean + over the hc streams; an unlisted layer writes nothing.""" meta = _make_metadata() + meta.captured_hidden_states.zero_() hs = torch.randn(4, HIDDEN, device="cuda", dtype=torch.bfloat16) meta.maybe_capture_hidden_states(59, hs) # layer 59 -> capture index 1 - got = meta.get_hidden_states(4) - assert torch.equal(got[:, HIDDEN : 2 * HIDDEN], hs) - - -def test_metadata_capture_hc_mean_reduction(): - """A flattened mHC residual [N, hc_mult*hidden] is reduced by mean over hc.""" - meta = _make_metadata() mhc = torch.randn(4, HC_MULT * HIDDEN, device="cuda", dtype=torch.bfloat16) meta.maybe_capture_hidden_states(58, mhc) # layer 58 -> capture index 0 - expected = mhc.reshape(4, HC_MULT, HIDDEN).mean(dim=1) + meta.maybe_capture_hidden_states(10, torch.randn(4, HIDDEN, device="cuda")) # unlisted got = meta.get_hidden_states(4) + assert torch.equal(got[:, HIDDEN : 2 * HIDDEN], hs) + expected = mhc.reshape(4, HC_MULT, HIDDEN).mean(dim=1) assert torch.equal(got[:, 0:HIDDEN], expected) + # The unlisted layer landed nowhere: the remaining slice stays zero. + assert torch.count_nonzero(got[:, 2 * HIDDEN :]) == 0 -def test_metadata_no_capture_for_unlisted_layer(): - meta = _make_metadata() - meta.captured_hidden_states.zero_() - meta.maybe_capture_hidden_states(10, torch.randn(4, HIDDEN, device="cuda")) - assert torch.count_nonzero(meta.get_hidden_states(4)) == 0 - - -def test_metadata_prepare_batch_indices(): - meta = _make_metadata() - meta.request_ids = [7, 3, 5] - meta.prepare() - assert meta.batch_indices_cuda[:3].tolist() == [0, 1, 2] - - -def _make_worker(): +def _make_worker(enable_confidence_scheduling=False, verify_len_tiers=None): cfg = types.SimpleNamespace( max_draft_len=5, spec_dec_mode=SpeculativeDecodingMode.DSPARK, - confidence_threshold=0.5, + enable_confidence_scheduling=enable_confidence_scheduling, + enable_ragged_verify=False, + verify_len_tiers=verify_len_tiers or [1, 3, 5], + confidence_sps_table_path=None, + confidence_sts_path=None, ) from tensorrt_llm.mapping import Mapping @@ -133,14 +127,11 @@ def test_worker_lazy_init_window_buffers(): worker._lazy_init(dm, meta) assert id(worker._kv_windows) == buf_id - -def test_worker_rejects_mismatched_block_size(): - worker = _make_worker() - draft_model = _fake_draft_model() - draft_model.block_size = 4 - + # A draft model whose block_size disagrees with the worker is refused. + bad = _fake_draft_model() + bad.block_size = 4 with pytest.raises(ValueError, match="block_size must equal worker max_draft_len"): - worker._lazy_init(draft_model, _make_metadata()) + _make_worker()._lazy_init(bad, _make_metadata()) def test_worker_slot_assignment_and_reset(): @@ -225,7 +216,8 @@ def write_context_windows(self, hidden, positions, windows): def test_prepare_builds_batch_to_slot_on_batched_path(): - """prepare() mirrors the host slot map into _batch_to_slot (default batched path).""" + """prepare() mirrors the host slot map into _batch_to_slot (default batched + path) and returns a dropped request's slot to the free pool.""" worker = _make_worker() meta = _make_metadata(max_num_requests=4) worker._lazy_init(_fake_draft_model(), meta) # batched is the default @@ -234,24 +226,13 @@ def test_prepare_builds_batch_to_slot_on_batched_path(): # Assign slots for two requests (as the prefill path would). sa = worker._assign_slot(100, reset=True) sb = worker._assign_slot(101, reset=True) + worker._ctx_len[sa] = 17 meta.request_ids = [101, 100] meta.prepare() # Mirror reflects request-order -> slot. assert worker._batch_to_slot[:2].tolist() == [sb, sa] - -def test_prepare_frees_stale_slots_on_batched_path(): - """A request that drops out of the batch returns its slot to the free pool.""" - worker = _make_worker() - meta = _make_metadata(max_num_requests=4) - worker._lazy_init(_fake_draft_model(), meta) - meta._dspark_worker = worker - - sa = worker._assign_slot(100, reset=True) - worker._assign_slot(101, reset=True) - worker._ctx_len[sa] = 17 - # Only request 101 survives; 100's slot must be freed + cleared. meta.request_ids = [101] meta.prepare() @@ -454,3 +435,137 @@ def fake_sample_draft_tokens(gl, sm, bs, *, num_contexts): # Verified tokens are surfaced unchanged. assert torch.equal(out["new_tokens"], accepted) assert torch.equal(out["new_tokens_lens"], num_accepted) + + +def _stride_probe_draft_model(block_size=5, num_stages=3, window_size=128, head_dim=64): + """Draft model stub that records the ``main_hidden`` it was handed.""" + seen = {} + + def forward_batched(main_hidden, bonus, start_pos, **kwargs): + seen["main_hidden"] = main_hidden.clone() + g = main_hidden.shape[0] + return ( + torch.zeros(g, block_size, dtype=torch.int32, device="cuda"), + None, + torch.zeros(g, block_size, 8, device="cuda"), + ) + + dm = types.SimpleNamespace( + num_stages=num_stages, + block_size=block_size, + _attn_params={"window_size": window_size, "head_dim": head_dim}, + write_context_windows_batched=lambda *a, **k: None, + forward_batched=forward_batched, + ) + return dm, seen + + +@pytest.mark.parametrize("runtime_draft_len", [3, None]) +def test_gen_draft_gathers_hidden_with_the_runtime_stride(runtime_draft_len): + """The gather stride must follow runtime_draft_len, not max_draft_len. + + The target lays out ``runtime_draft_len + 1`` tokens per generation request. + Striding by ``max_draft_len + 1`` instead walks into the *next* request's + hidden states as soon as the scheduler trims the draft length -- the draft + then conditions on another request's context, with no error anywhere. + ``runtime_draft_len=None`` pins the fallback: metadata without a runtime + length keeps the pre-existing max_draft_len stride. + """ + num_gens, hidden = 4, HIDDEN * NCAP + worker = _make_worker() + dm, seen = _stride_probe_draft_model() + meta = _make_metadata(max_num_requests=8, max_num_tokens=256) + worker._lazy_init(dm, meta) + worker._batch_to_slot[:num_gens] = torch.arange(num_gens, device="cuda") + + # Row r of request g is tagged with g so a cross-request read is visible. + stride = (runtime_draft_len or worker.max_draft_len) + 1 + captured = torch.zeros(num_gens * stride, hidden, device="cuda") + for g in range(num_gens): + captured[g * stride : (g + 1) * stride] = float(g) + + meta.runtime_draft_len = runtime_draft_len + meta.get_hidden_states = lambda _n: captured + attn = types.SimpleNamespace(num_ctx_tokens=0, num_seqs=num_gens, num_contexts=0) + + # Everyone accepts exactly one token, so the bonus sits at offset 0 of each + # request's block: main_hidden[g] must be exactly g. + nacc = torch.ones(num_gens, dtype=torch.int32, device="cuda") + accepted = torch.zeros(num_gens, worker.max_draft_len + 1, dtype=torch.int32, device="cuda") + + worker._draft_gen_block_batched( + dm, + meta, + attn, + accepted, + nacc, + num_contexts=0, + batch_size=num_gens, + total_target_tokens=captured.shape[0], + ) + + got = seen["main_hidden"][:, 0].tolist() + assert got == [float(g) for g in range(num_gens)], ( + f"runtime_draft_len={runtime_draft_len}: drafted from the wrong request's " + f"hidden states, got {got}" + ) + + +def test_confidence_rows_are_looked_up_by_slot(): + """The confidence buffer is slot-indexed; a batch rarely owns slots 0..G-1. + + The snapshot the scheduler reads is one iteration old, and joins/departures + reshuffle the batch in between, so batch position is not a usable key. + ``confidence_row_for`` re-associates each request with the row its own draft + wrote. A request with no slot yet falls back to a dedicated, permanently + neutral row ("verify the whole block") past every real slot -- the scratch + row is not safe for that, since padded and unknown requests write through it. + """ + worker = _make_worker(enable_confidence_scheduling=True) + dm, _ = _stride_probe_draft_model() + worker._lazy_init(dm, _make_metadata(max_num_requests=8)) + assert worker._confidence_logits is not None + + # Unscored request: the permanently-neutral row, past every real slot AND + # past the scratch row, and it reads as ~certain acceptance. + neutral = worker.confidence_row_for(12345) + assert neutral == worker._neutral_conf_row + assert neutral > worker._scratch_slot + buf = worker.staged_confidence_buffer() + assert torch.sigmoid(buf[neutral]).min().item() > 0.999 + + # Tag every real row with its index, then hand three requests explicit slots. + for s in range(worker._scratch_slot): + worker._confidence_logits[s] = float(s) + worker._req_to_slot = {100: 5, 101: 2, 102: 7} + + rows = [worker.confidence_row_for(r) for r in (100, 101, 102)] + assert rows == [5, 2, 7] + assert buf[rows][:, 0].tolist() == [5.0, 2.0, 7.0] + + +def test_a_freshly_assigned_slot_starts_neutral(): + """Otherwise a new request's first decision trims on a dead draft's scores. + + Slots are recycled between requests, so whatever the previous occupant's + draft scored is still sitting in the row when the next request moves in. + """ + worker = _make_worker(enable_confidence_scheduling=True) + dm, _ = _stride_probe_draft_model() + worker._lazy_init(dm, _make_metadata(max_num_requests=8)) + + # Every real slot carries a dead previous occupant's hopeless confidence. + worker._confidence_logits[: worker._scratch_slot] = -20.0 + + slot = worker._assign_slot(100, reset=False) + assert slot < worker._scratch_slot + assert torch.sigmoid(worker._confidence_logits[slot]).min().item() > 0.999 + + +def test_confidence_buffer_absent_when_scheduling_is_off(): + """Zero cost when the feature is off.""" + worker = _make_worker(enable_confidence_scheduling=False) + dm, _ = _stride_probe_draft_model() + worker._lazy_init(dm, _make_metadata(max_num_requests=8)) + assert worker._confidence_logits is None + assert worker.verify_planner is None diff --git a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py deleted file mode 100644 index fbfaca610386..000000000000 --- a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py +++ /dev/null @@ -1,394 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Unit tests for the synthetic forced acceptance rate in 1-model spec dec. - -These tests exercise ``SpecWorkerBase._apply_force_accepted_tokens``, which -honors the ``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` environment -variable. The implementation must be: - - - Faithful to fractional rates: integer part always accepts, fractional - part acts as the per-iteration probability of accepting one extra draft - token. - - Tensor-parallel deterministic: identical seed/counter state across - ranks so all ranks agree on per-request accepted counts (otherwise - downstream collectives expecting matching shapes deadlock). - - CUDA-graph compatible: pure-tensor ops with statically-known shapes, - captureable and replayable. -""" - -import os -from typing import Optional -from unittest.mock import patch - -import pytest -import torch - -from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata -from tensorrt_llm._torch.speculative.interface import ( - FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR, - SpecWorkerBase, - get_force_num_accepted_tokens, - get_force_num_accepted_tokens_float, -) - - -class _StubSpecWorker(SpecWorkerBase): - """Concrete ``SpecWorkerBase`` that stubs out the abstract API. - - Used purely to drive ``_apply_force_accepted_tokens`` in isolation. - """ - - @property - def max_draft_len(self) -> int: - return 8 - - def _forward_impl(self, *args: object, **kwargs: object) -> None: - raise NotImplementedError - - -def _make_worker(value: Optional[float] = None) -> _StubSpecWorker: - worker = _StubSpecWorker() - if value is not None: - worker.force_num_accepted_tokens = value - return worker - - -def _make_input(batch_size: int, device: str = "cuda") -> torch.Tensor: - # Initial counts: 1 (target token, no draft accepts yet). The forced - # override only ever rewrites entries from ``num_contexts:`` onward. - return torch.ones(batch_size, dtype=torch.int32, device=device) - - -def _require_cuda(): - if not torch.cuda.is_available(): - pytest.skip("CUDA required for this test.") - - -# ---------------- env-var parsing helpers ----------------------------------- - - -@pytest.mark.parametrize( - "env_value, expected", - [ - ("0", 0.0), - ("3", 3.0), - ("2.6", 2.6), - ("0.5", 0.5), - ("not-a-number", 0.0), - ], -) -def test_get_force_num_accepted_tokens_float(env_value, expected): - with patch.dict(os.environ, {FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR: env_value}): - assert get_force_num_accepted_tokens_float() == pytest.approx(expected) - - -@pytest.mark.parametrize( - "env_value, expected", - [ - ("0", 0), - ("3", 3), - ("2.6", 0), - ], -) -def test_get_force_num_accepted_tokens_int_unchanged(env_value, expected): - """The int helper must keep its original behavior (used by 2-model).""" - with patch.dict(os.environ, {FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR: env_value}): - assert get_force_num_accepted_tokens() == expected - - -# ---------------- _apply_force_accepted_tokens semantics -------------------- - - -def test_zero_value_is_noop(): - _require_cuda() - worker = _make_worker(0.0) - base = _make_input(batch_size=4) - before = base.clone() - out = worker._apply_force_accepted_tokens(base, num_contexts=0, runtime_draft_len=4) - assert torch.equal(out, before) - # No RNG state should have been touched in the early-exit path. - assert worker._force_accept_rng_pool is None - assert worker._force_accept_rng_counter is None - - -@pytest.mark.parametrize("value", [1.0, 2.0, 4.0]) -def test_integer_value_matches_legacy_behavior(value): - _require_cuda() - worker = _make_worker(value) - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=8), num_contexts=0, runtime_draft_len=4 - ) - expected = min(int(value) + 1, 4 + 1) - assert torch.all(out == expected) - # Pure-integer path must not allocate the RNG pool. - assert worker._force_accept_rng_pool is None - - -def test_integer_value_caps_at_runtime_draft_len_plus_one(): - _require_cuda() - worker = _make_worker(10.0) - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=4), num_contexts=0, runtime_draft_len=2 - ) - # 10 draft tokens requested but only 2 available → all 3 (= 2 + target). - assert torch.all(out == 3) - - -def test_fractional_only_emits_two_values(): - _require_cuda() - worker = _make_worker(2.6) - seen = set() - for _ in range(50): - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=64), num_contexts=0, runtime_draft_len=4 - ) - seen.update(out.unique().tolist()) - # Either 2 draft + target = 3, or 3 draft + target = 4. - assert seen == {3, 4} - - -def test_fractional_distribution_matches_target_probability(): - """Average over many draws should converge to the target fraction.""" - _require_cuda() - worker = _make_worker(2.6) - target_frac = 0.6 - n_iters = 200 - batch = 64 - extra_count = 0 - total = 0 - for _ in range(n_iters): - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=batch), num_contexts=0, runtime_draft_len=4 - ) - extra_count += int((out == 4).sum().item()) - total += batch - measured = extra_count / total - # Pool-based RNG is deterministic; the empirical mean over ~13k draws - # is tight against 0.6. - assert abs(measured - target_frac) < 0.03, ( - f"measured fraction {measured:.4f} differs from target {target_frac}" - ) - - -def test_fractional_capped_when_no_room(): - """If ``int_part + 1`` already saturates the cap, no extra is granted.""" - _require_cuda() - worker = _make_worker(9.5) - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=8), num_contexts=0, runtime_draft_len=2 - ) - # max_total = runtime_draft_len + 1 = 3, base_total = min(10, 3) = 3. - # frac_part > 0 but ``base_total < max_total`` is False → all 3, no RNG. - assert torch.all(out == 3) - assert worker._force_accept_rng_pool is None - - -def test_num_contexts_offset_is_respected(): - """Context (prefill) rows must be left untouched by the override.""" - _require_cuda() - worker = _make_worker(2.6) - batch_size = 8 - num_contexts = 3 - base = _make_input(batch_size=batch_size) - sentinel = torch.iinfo(base.dtype).max - base[:num_contexts] = sentinel - out = worker._apply_force_accepted_tokens(base, num_contexts=num_contexts, runtime_draft_len=4) - assert torch.all(out[:num_contexts] == sentinel) - assert torch.all((out[num_contexts:] == 3) | (out[num_contexts:] == 4)) - - -# ---------------- TP determinism -------------------------------------------- - - -def test_tp_determinism_across_independent_workers(): - """ - Two independently-instantiated workers (simulating two TP ranks) must - produce bit-identical accepted-token counts when fed the same call - sequence. This is the property whose absence used to hang TP>1. - """ - _require_cuda() - rank0 = _make_worker(2.6) - rank1 = _make_worker(2.6) - # Use a non-power-of-two batch to make accidental shape coincidences - # less likely to mask divergence. - for _ in range(32): - out0 = rank0._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - out1 = rank1._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - assert torch.equal(out0, out1) - - -def test_tp_determinism_survives_default_generator_drift(): - """ - Even if the default CUDA generator state diverges between ranks (e.g. - because some other component consumed random numbers), the synthetic - AR draws must still match. This is the regression guard for the - original bug, where ``torch.rand`` against the default generator gave - each rank a different sequence. - """ - _require_cuda() - rank0 = _make_worker(2.6) - rank1 = _make_worker(2.6) - - def _advance_default_generator(seed: int): - # Seed the default CUDA generator on the same device differently - # for each "rank" and consume some randoms, simulating drift. - torch.cuda.manual_seed(seed) - _ = torch.rand(128, device="cuda") - - _advance_default_generator(seed=11) - out0 = rank0._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - _advance_default_generator(seed=999) - out1 = rank1._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - assert torch.equal(out0, out1) - - -# ---------------- CUDA graph capture/replay --------------------------------- - - -def test_cuda_graph_capture_and_replay_match_eager(): - """ - Capturing the override into a CUDA graph and replaying it must produce - the exact same per-replay output sequence as running it eagerly. - """ - _require_cuda() - - n_iters = 8 - batch_size = 16 - runtime_draft_len = 4 - value = 2.6 - - # 1) Eager reference. - eager_worker = _make_worker(value) - eager_outputs = [] - for _ in range(n_iters): - out = eager_worker._apply_force_accepted_tokens( - _make_input(batch_size=batch_size), num_contexts=0, runtime_draft_len=runtime_draft_len - ) - eager_outputs.append(out.clone()) - - # 2) Captured-graph run, aligned to the eager reference. - graph_worker = _make_worker(value) - static_input = _make_input(batch_size=batch_size) - # Eager warmup: forces lazy RNG state allocation OUTSIDE capture. - graph_worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - # Realign the device-side counter so the captured graph's first replay - # advances 0 → 1, matching the eager loop's first iteration. - graph_worker._force_accept_rng_counter.zero_() - - static_input.fill_(1) - static_output = torch.empty_like(static_input) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - # ``_apply_force_accepted_tokens`` mutates ``static_input`` in place; - # mirror it into ``static_output`` so we can snapshot per replay. - static_output.copy_(static_input) - - graph_outputs = [] - for _ in range(n_iters): - # Reset input each replay so the override semantics match eager - # (which also starts from all-1s every iteration). - static_input.fill_(1) - graph.replay() - torch.cuda.synchronize() - graph_outputs.append(static_output.clone()) - - for i, (eager, graphed) in enumerate(zip(eager_outputs, graph_outputs)): - assert torch.equal(eager, graphed), ( - f"Iteration {i}: eager={eager.tolist()} vs graphed={graphed.tolist()}" - ) - - -def test_cuda_graph_replay_advances_rng_state(): - """ - Across replays, the captured counter must advance, producing different - (but deterministic) draws. A constant output across replays would mean - the counter increment was not captured. - """ - _require_cuda() - - worker = _make_worker(0.5) - batch_size = 64 - runtime_draft_len = 4 - - static_input = _make_input(batch_size=batch_size) - worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - worker._force_accept_rng_counter.zero_() - - static_input.fill_(1) - static_output = torch.empty_like(static_input) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - static_output.copy_(static_input) - - snapshots = [] - for _ in range(8): - static_input.fill_(1) - graph.replay() - torch.cuda.synchronize() - snapshots.append(static_output.clone()) - - # At least two replays must differ; otherwise the captured counter is - # not advancing (the original bug-class). - assert any(not torch.equal(snapshots[0], s) for s in snapshots[1:]), ( - "Captured graph produced identical outputs on every replay — " - "RNG counter is not advancing inside the captured graph." - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) - - -class _FailingSpecWorker(_StubSpecWorker): - """Stub whose forward saves spec-dec metadata state and then fails.""" - - def _forward_impl(self, *args: object, **kwargs: object) -> None: - attn_metadata = kwargs["attn_metadata"] - attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") - raise RuntimeError("simulated draft failure") - - -def test_forward_restores_spec_dec_state_on_failure() -> None: - """A failure between prepare_for_spec_dec and restore must not leak saved - attn-metadata state; SpecWorkerBase.forward restores it in its cleanup - (https://nvbugs/6442074).""" - _require_cuda() - attn_metadata = AttentionMetadata(max_num_requests=2, max_num_tokens=16) - attn_metadata.seq_lens = torch.ones(2, dtype=torch.int32) - worker = _FailingSpecWorker() - for _ in range(2): - # The second iteration would trip the pairing assert inside - # prepare_for_spec_dec if the first failure had leaked saved state. - with pytest.raises(RuntimeError, match="simulated draft failure"): - worker(attn_metadata=attn_metadata, spec_metadata=None) - assert not attn_metadata.has_spec_dec_saved_state diff --git a/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py b/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py deleted file mode 100644 index bdcc3fd77992..000000000000 --- a/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py +++ /dev/null @@ -1,88 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""CPU unit tests for the group-synchronized ``is_all_greedy_sample`` override. - -Under ADP + LM-head TP with rejection sampling, the greedy-vs-advanced path -choice gates group collectives, so the model engine all-gathers the per-rank -flags and stores the group AND in ``SpecMetadata.group_all_greedy_sample``; -``_scan_one_model_sampling`` must then re-apply it on every rescan (populate -runs after the CUDA graph key is built and would otherwise resurrect the -rank-local value). - -These tests call ``_scan_one_model_sampling`` unbound on a SimpleNamespace -stand-in, mirroring test_rejection_buffers_guard.py, so no GPU or full -SpecMetadata construction is needed. -""" - -import types - -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.speculative.interface import SpecMetadata - - -def _fake_request(temperature=None, top_k=None, top_p=None, slot=0): - return types.SimpleNamespace( - sampling_config=types.SimpleNamespace( - temperature=[temperature] if temperature is not None else None, - top_k=[top_k] if top_k is not None else None, - top_p=[top_p] if top_p is not None else None, - ), - state=LlmRequestState.GENERATION_IN_PROGRESS, - py_seq_slot=slot, - ) - - -def _fake_meta(group_all_greedy_sample=None, force_capture=False): - return types.SimpleNamespace( - runtime_draft_len=2, - dummy_slot_row=0, - group_all_greedy_sample=group_all_greedy_sample, - _force_non_greedy_for_capture=force_capture, - ) - - -def _scan(meta, requests): - return SpecMetadata._scan_one_model_sampling(meta, requests) - - -def test_local_value_used_when_no_group_sync(): - meta = _fake_meta(group_all_greedy_sample=None) - _scan(meta, [_fake_request(), _fake_request()]) - assert meta.is_all_greedy_sample is True - - _scan(meta, [_fake_request(), _fake_request(temperature=0.8)]) - assert meta.is_all_greedy_sample is False - - -def test_group_override_pulls_greedy_rank_onto_advanced_path(): - # This rank's batch is all-greedy, but another rank in the LM-head-TP - # group has a sampling request: the group AND (False) must win so the - # whole group takes the advanced path together. - meta = _fake_meta(group_all_greedy_sample=False) - _scan(meta, [_fake_request(), _fake_request()]) - assert meta.is_all_greedy_sample is False - - -def test_group_override_survives_rescan(): - # populate_sampling_params_for_one_model rescans after the CUDA graph key - # is built; the override must keep applying so the key, the buffers, and - # the worker branches all agree. - meta = _fake_meta(group_all_greedy_sample=False) - for _ in range(3): - _scan(meta, [_fake_request()]) - assert meta.is_all_greedy_sample is False - - -def test_group_override_true_keeps_greedy(): - meta = _fake_meta(group_all_greedy_sample=True) - _scan(meta, [_fake_request()]) - assert meta.is_all_greedy_sample is True - - -def test_capture_override_composes_with_group_sync(): - # Warmup forces the advanced variant to capture its CUDA graph; the group - # value is derived from capture-forced locals (all False), so the final - # flag stays False regardless of composition order. - meta = _fake_meta(group_all_greedy_sample=False, force_capture=True) - _scan(meta, [_fake_request()]) - assert meta.is_all_greedy_sample is False diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py deleted file mode 100644 index a71b9f1e6640..000000000000 --- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py +++ /dev/null @@ -1,275 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU unit tests for the one-model rejection slot-buffer allocation and the -fail-closed acceptance guard. - -These exercise the pure buffer/guard logic with synthetic tensors and lightweight -attribute stand-ins (``types.SimpleNamespace``), so they do not construct a full -``SpecMetadata`` or run any model forward: - -- ``SpecMetadata.prepare_rejection_sampling_buffers`` only reads/writes ``self`` - attributes (``use_rejection_sampling``, ``vocab_size``, ``max_num_requests``, - ``max_draft_len``, ``draft_probs``, ``batch_slot_ids``, ``full_draft_probs``). -- ``SpecWorkerBase._rejection_buffers_valid`` only reads its arguments and - ``spec_metadata`` attributes (no ``self`` use), so it can be called unbound. - -Requires CUDA because the buffers are allocated on ``device='cuda'``. -""" - -import types - -import pytest -import torch - -from tensorrt_llm._torch.speculative.interface import SpecMetadata, SpecWorkerBase - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), reason="rejection buffers are CUDA tensors" -) - -R, K, V = 8, 4, 32 # max_num_requests, max_draft_len, vocab_size - - -def _alloc_meta(**over): - base = dict( - use_rejection_sampling=True, - draft_probs=None, - vocab_size=V, - draft_vocab_size=V, - max_num_requests=R, - max_draft_len=K, - num_seq_slots=0, - batch_slot_ids=None, - full_draft_probs=None, - draft_probs_vocab_size=0, - ) - base.update(over) - return types.SimpleNamespace(**base) - - -def test_prepare_buffers_allocates_when_enabled(): - # Shared draft/target vocab: draft_probs + batch_slot_ids are allocated; - # full_draft_probs is skipped (only needed when the vocabs differ). - m = _alloc_meta() - SpecMetadata.prepare_rejection_sampling_buffers(m) - assert m.draft_probs is not None and tuple(m.draft_probs.shape) == (R + 1, K, V) - assert m.batch_slot_ids is not None and m.batch_slot_ids.shape[0] == R - assert m.batch_slot_ids.dtype == torch.long - assert m.full_draft_probs is None - assert m.draft_probs_vocab_size == V - - -def test_prepare_buffers_allocates_full_draft_probs_on_vocab_mismatch(): - # Distinct draft vocab: full_draft_probs (d2t-expanded) is allocated. - m = _alloc_meta(draft_vocab_size=V - 1) - SpecMetadata.prepare_rejection_sampling_buffers(m) - assert m.full_draft_probs is not None and tuple(m.full_draft_probs.shape) == (R + 1, K, V) - - -def test_prepare_buffers_span_seq_slot_pool(): - # Under DeepSeek-V4 overlap scheduling the SeqSlotManager pool - # (num_seq_slots) can exceed max_num_requests; py_seq_slot then indexes past - # max_num_requests. Slot-indexed buffers must span the full pool plus the - # dummy scratch row, and dummy_slot_row must land on that last row so a real - # slot never aliases the dummy row. - pool = 2 * R - m = _alloc_meta(num_seq_slots=pool) - SpecMetadata.prepare_rejection_sampling_buffers(m) - assert tuple(m.draft_probs.shape) == (pool + 1, K, V) - assert m.dummy_slot_row == pool - - -def test_prepare_buffers_noop_when_disabled(): - m = _alloc_meta(use_rejection_sampling=False) - SpecMetadata.prepare_rejection_sampling_buffers(m) - assert m.draft_probs is None - assert m.batch_slot_ids is None - assert m.full_draft_probs is None - - -def _valid_state(): - return types.SimpleNamespace( - draft_probs=torch.empty((R, K, V), device="cuda"), - batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), - ) - - -# num_contexts=0, num_gens=4 -_NUM_CTX = 0 -_BATCH = 4 -_NUM_GENS = _BATCH - _NUM_CTX - - -def _good_args(): - draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") - logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") - return draft_tokens, K, V, _NUM_CTX, _BATCH, logits - - -_DUMMY = object() # _rejection_buffers_valid does not use self - - -def _call(meta, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits): - return SpecWorkerBase._rejection_buffers_valid( - _DUMMY, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits, meta - ) - - -def test_guard_true_on_valid_state(): - assert _call(_valid_state(), *_good_args()) is True - - -# Each builder returns (meta, *_call args) for one malformed state that the -# fail-closed guard must reject. Kept as a single parametrized test so the guard's -# rejection branches read as one table. -def _case_draft_probs_missing(): - m = _valid_state() - m.draft_probs = None - return (m, *_good_args()) - - -def _case_batch_slot_ids_missing(): - m = _valid_state() - m.batch_slot_ids = None - return (m, *_good_args()) - - -def _case_stored_vocab_nonpositive(): - dt, dl, _sv, nc, bs, lg = _good_args() - return (_valid_state(), dt, dl, 0, nc, bs, lg) - - -def _case_stored_vocab_exceeds_buffer(): - dt, dl, _sv, nc, bs, lg = _good_args() - return (_valid_state(), dt, dl, V + 1, nc, bs, lg) - - -def _case_draft_len_exceeds_buffer(): - # draft_probs has K steps; ask for K+1. - draft_tokens = torch.zeros((_NUM_GENS, K + 1), dtype=torch.int, device="cuda") - logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 2), V), device="cuda") - return (_valid_state(), draft_tokens, K + 1, V, _NUM_CTX, _BATCH, logits) - - -def _case_draft_tokens_wrong_rows(): - dt = torch.zeros((_NUM_GENS + 1, K), dtype=torch.int, device="cuda") - _, dl, sv, nc, bs, lg = _good_args() - return (_valid_state(), dt, dl, sv, nc, bs, lg) - - -def _case_too_few_logits_rows(): - dt, dl, sv, nc, bs, _lg = _good_args() - too_few = torch.zeros((nc + _NUM_GENS, V), device="cuda") # missing +1 each - return (_valid_state(), dt, dl, sv, nc, bs, too_few) - - -def _case_batch_slot_ids_too_short(): - m = _valid_state() - m.batch_slot_ids = torch.arange(_BATCH - 1, device="cuda", dtype=torch.long) - return (m, *_good_args()) - - -@pytest.mark.parametrize( - "build_case", - [ - _case_draft_probs_missing, - _case_batch_slot_ids_missing, - _case_stored_vocab_nonpositive, - _case_stored_vocab_exceeds_buffer, - _case_draft_len_exceeds_buffer, - _case_draft_tokens_wrong_rows, - _case_too_few_logits_rows, - _case_batch_slot_ids_too_short, - ], - ids=lambda f: f.__name__.removeprefix("_case_"), -) -def test_guard_false_on_malformed_state(build_case): - assert _call(*build_case()) is False - - -# -------------------------------------------------------------------------- -# Fail-closed acceptance dispatch: prove _accept_draft_tokens() routes to -# strict/base acceptance when the buffers are malformed, and routes to the -# rejection method when the state is valid. The two acceptance methods are -# stubbed to record which path ran; the real _can_use_rejection_sampling and -# _rejection_buffers_valid are exercised. -# -------------------------------------------------------------------------- - - -class _Worker(SpecWorkerBase): - @property - def max_draft_len(self) -> int: - return K - - def _forward_impl(self, *args: object, **kwargs: object) -> None: - raise NotImplementedError - - -def _dispatch_meta(**over): - base = dict( - use_rejection_sampling=True, - is_all_greedy_sample=False, - draft_probs_vocab_size=V, - draft_probs_last_dim=V, - batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), - draft_probs=torch.zeros((R, K, V), device="cuda"), - ) - base.update(over) - return types.SimpleNamespace(**base) - - -def _make_worker(): - w = _Worker() - calls = {"base": 0, "rejection": 0} - - def _base(logits, draft_tokens, num_contexts, batch_size, spec_metadata): - calls["base"] += 1 - return ("base", None) - - def _rej(logits, draft_tokens, draft_probs, num_contexts, batch_size, spec_metadata): - calls["rejection"] += 1 - return ("rejection", None) - - w._sample_and_accept_draft_tokens_base = _base - w._sample_and_accept_draft_tokens_rejection = _rej - return w, calls - - -def test_accept_dispatch_fails_closed_on_malformed_buffers(): - w, calls = _make_worker() - # Buffers are malformed (draft_probs None) -> _rejection_buffers_valid must - # return False, so acceptance falls back to base and the rejection kernel is - # skipped. - meta = _dispatch_meta(draft_probs=None) - draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") - logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") - out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) - assert out == ("base", None) - assert calls["base"] == 1 and calls["rejection"] == 0 - - -def test_accept_dispatch_routes_to_rejection_on_valid_state(): - w, calls = _make_worker() - meta = _dispatch_meta() # valid buffers - draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") - logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") - out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) - assert out == ("rejection", None) - assert calls["rejection"] == 1 and calls["base"] == 0 - - -def test_accept_dispatch_base_when_all_greedy(): - w, calls = _make_worker() - # All-greedy batch: rejection is bypassed regardless of buffers. - meta = _dispatch_meta(is_all_greedy_sample=True) - draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") - logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") - out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) - assert out == ("base", None) - assert calls["base"] == 1 and calls["rejection"] == 0 - - -if __name__ == "__main__": - import sys - - sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/unittest/_torch/speculative/test_spec_sampler.py b/tests/unittest/_torch/speculative/test_spec_sampler.py new file mode 100644 index 000000000000..055a0bd981c7 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_spec_sampler.py @@ -0,0 +1,650 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the one-model spec-decode sampler/acceptance framework: +forced acceptance (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS``), +rejection-sampling buffer allocation with its fail-closed guard, and the +group-synchronized ``is_all_greedy_sample`` override.""" + +import os +import types +from typing import Optional +from unittest.mock import patch + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.speculative.interface import ( + FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR, + SpecMetadata, + SpecWorkerBase, + get_force_num_accepted_tokens, + get_force_num_accepted_tokens_float, +) + +_cuda_only = pytest.mark.skipif( + not torch.cuda.is_available(), reason="rejection buffers are CUDA tensors" +) + + +class _StubSpecWorker(SpecWorkerBase): + """Concrete ``SpecWorkerBase`` that stubs out the abstract API.""" + + @property + def max_draft_len(self) -> int: + return 8 + + def _forward_impl(self, *args: object, **kwargs: object) -> None: + raise NotImplementedError + + +def _make_worker(value: Optional[float] = None) -> _StubSpecWorker: + worker = _StubSpecWorker() + if value is not None: + worker.force_num_accepted_tokens = value + return worker + + +def _make_input(batch_size: int, device: str = "cuda") -> torch.Tensor: + # Initial counts: 1 (target token, no draft accepts yet). The forced + # override only ever rewrites entries from ``num_contexts:`` onward. + return torch.ones(batch_size, dtype=torch.int32, device=device) + + +def _require_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required for this test.") + + +# ---------------- env-var parsing helpers ----------------------------------- + + +@pytest.mark.parametrize( + "env_value, expected_float, expected_int", + [ + ("0", 0.0, 0), + ("3", 3.0, 3), + ("2.6", 2.6, 0), + ("0.5", 0.5, 0), + ("not-a-number", 0.0, 0), + ], +) +def test_get_force_num_accepted_tokens_float(env_value, expected_float, expected_int): + """The float helper parses fractional rates; the int helper (used by the + 2-model path) must keep its original behavior on the same values.""" + with patch.dict(os.environ, {FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR: env_value}): + assert get_force_num_accepted_tokens_float() == pytest.approx(expected_float) + assert get_force_num_accepted_tokens() == expected_int + + +# ---------------- _apply_force_accepted_tokens semantics -------------------- + + +def test_zero_value_is_noop(): + """A zero forced rate must leave inputs and RNG state untouched.""" + _require_cuda() + worker = _make_worker(0.0) + base = _make_input(batch_size=4) + before = base.clone() + out = worker._apply_force_accepted_tokens(base, num_contexts=0, runtime_draft_len=4) + assert torch.equal(out, before) + # No RNG state should have been touched in the early-exit path. + assert worker._force_accept_rng_pool is None + assert worker._force_accept_rng_counter is None + + +@pytest.mark.parametrize( + "value, runtime_draft_len", + [ + (1.0, 4), + (2.0, 4), + (4.0, 4), + # 10 draft tokens requested but only 2 available -> capped at 3. + (10.0, 2), + # Fractional but no room: max_total = 3, base_total = min(10, 3) = 3, + # so ``base_total < max_total`` is False -> all 3, no RNG draw. + (9.5, 2), + ], +) +def test_integer_value_matches_legacy_behavior(value, runtime_draft_len): + """Integer (or capped) rates take the deterministic path with no RNG pool.""" + _require_cuda() + worker = _make_worker(value) + out = worker._apply_force_accepted_tokens( + _make_input(batch_size=8), num_contexts=0, runtime_draft_len=runtime_draft_len + ) + expected = min(int(value) + 1, runtime_draft_len + 1) + assert torch.all(out == expected) + # The deterministic (non-random) path must not allocate the RNG pool. + assert worker._force_accept_rng_pool is None + + +def test_fractional_distribution_matches_target_probability(): + """Fractional draws emit only base/base+1 and average to the target rate.""" + _require_cuda() + worker = _make_worker(2.6) + target_frac = 0.6 + n_iters = 200 + batch = 64 + seen = set() + extra_count = 0 + total = 0 + for _ in range(n_iters): + out = worker._apply_force_accepted_tokens( + _make_input(batch_size=batch), num_contexts=0, runtime_draft_len=4 + ) + seen.update(out.unique().tolist()) + extra_count += int((out == 4).sum().item()) + total += batch + # Either 2 draft + target = 3, or 3 draft + target = 4 -- never anything else. + assert seen == {3, 4} + measured = extra_count / total + assert abs(measured - target_frac) < 0.03, ( + f"measured fraction {measured:.4f} differs from target {target_frac}" + ) + + +def test_num_contexts_offset_is_respected(): + """Context (prefill) rows must be left untouched by the override.""" + _require_cuda() + worker = _make_worker(2.6) + batch_size = 8 + num_contexts = 3 + base = _make_input(batch_size=batch_size) + sentinel = torch.iinfo(base.dtype).max + base[:num_contexts] = sentinel + out = worker._apply_force_accepted_tokens(base, num_contexts=num_contexts, runtime_draft_len=4) + assert torch.all(out[:num_contexts] == sentinel) + assert torch.all((out[num_contexts:] == 3) | (out[num_contexts:] == 4)) + + +# ---------------- TP determinism -------------------------------------------- + + +def test_tp_determinism_across_independent_workers(): + """Two independently-instantiated workers (simulating two TP ranks) must + produce bit-identical accepted counts for the same call sequence.""" + _require_cuda() + rank0 = _make_worker(2.6) + rank1 = _make_worker(2.6) + # Use a non-power-of-two batch to make accidental shape coincidences + # less likely to mask divergence. + for _ in range(32): + out0 = rank0._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + out1 = rank1._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + assert torch.equal(out0, out1) + + +def test_tp_determinism_survives_default_generator_drift(): + """Synthetic AR draws must still match across ranks even when the default + CUDA generator state has diverged between them.""" + _require_cuda() + rank0 = _make_worker(2.6) + rank1 = _make_worker(2.6) + + def _advance_default_generator(seed: int): + # Seed the default CUDA generator on the same device differently + # for each "rank" and consume some randoms, simulating drift. + torch.cuda.manual_seed(seed) + _ = torch.rand(128, device="cuda") + + _advance_default_generator(seed=11) + out0 = rank0._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + _advance_default_generator(seed=999) + out1 = rank1._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + assert torch.equal(out0, out1) + + +# ---------------- CUDA graph capture/replay --------------------------------- + + +def test_cuda_graph_capture_and_replay_match_eager(): + """Capturing the override into a CUDA graph and replaying it must produce + the exact same per-replay output sequence as running it eagerly.""" + _require_cuda() + + n_iters = 8 + batch_size = 16 + runtime_draft_len = 4 + value = 2.6 + + # 1) Eager reference. + eager_worker = _make_worker(value) + eager_outputs = [] + for _ in range(n_iters): + out = eager_worker._apply_force_accepted_tokens( + _make_input(batch_size=batch_size), num_contexts=0, runtime_draft_len=runtime_draft_len + ) + eager_outputs.append(out.clone()) + + # 2) Captured-graph run, aligned to the eager reference. + graph_worker = _make_worker(value) + static_input = _make_input(batch_size=batch_size) + # Eager warmup: forces lazy RNG state allocation OUTSIDE capture. + graph_worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + # Realign the device-side counter so the captured graph's first replay + # advances 0 → 1, matching the eager loop's first iteration. + graph_worker._force_accept_rng_counter.zero_() + + static_input.fill_(1) + static_output = torch.empty_like(static_input) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + # ``_apply_force_accepted_tokens`` mutates ``static_input`` in place; + # mirror it into ``static_output`` so we can snapshot per replay. + static_output.copy_(static_input) + + graph_outputs = [] + for _ in range(n_iters): + # Reset input each replay so the override semantics match eager + # (which also starts from all-1s every iteration). + static_input.fill_(1) + graph.replay() + torch.cuda.synchronize() + graph_outputs.append(static_output.clone()) + + for i, (eager, graphed) in enumerate(zip(eager_outputs, graph_outputs)): + assert torch.equal(eager, graphed), ( + f"Iteration {i}: eager={eager.tolist()} vs graphed={graphed.tolist()}" + ) + + +def test_cuda_graph_replay_advances_rng_state(): + """Across replays the captured RNG counter must advance, producing + different (but deterministic) draws rather than a constant output.""" + _require_cuda() + + worker = _make_worker(0.5) + batch_size = 64 + runtime_draft_len = 4 + + static_input = _make_input(batch_size=batch_size) + worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + worker._force_accept_rng_counter.zero_() + + static_input.fill_(1) + static_output = torch.empty_like(static_input) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + static_output.copy_(static_input) + + snapshots = [] + for _ in range(8): + static_input.fill_(1) + graph.replay() + torch.cuda.synchronize() + snapshots.append(static_output.clone()) + + # At least two replays must differ; otherwise the captured counter is + # not advancing (the original bug-class). + assert any(not torch.equal(snapshots[0], s) for s in snapshots[1:]), ( + "Captured graph produced identical outputs on every replay — " + "RNG counter is not advancing inside the captured graph." + ) + + +class _FailingSpecWorker(_StubSpecWorker): + """Stub whose forward saves spec-dec metadata state and then fails.""" + + def _forward_impl(self, *args: object, **kwargs: object) -> None: + attn_metadata = kwargs["attn_metadata"] + attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") + raise RuntimeError("simulated draft failure") + + +def test_forward_restores_spec_dec_state_on_failure() -> None: + """A failure between prepare_for_spec_dec and restore must not leak saved + attn-metadata state; SpecWorkerBase.forward restores it in its cleanup.""" + _require_cuda() + attn_metadata = AttentionMetadata(max_num_requests=2, max_num_tokens=16) + attn_metadata.seq_lens = torch.ones(2, dtype=torch.int32) + worker = _FailingSpecWorker() + for _ in range(2): + # The second iteration would trip the pairing assert inside + # prepare_for_spec_dec if the first failure had leaked saved state. + with pytest.raises(RuntimeError, match="simulated draft failure"): + worker(attn_metadata=attn_metadata, spec_metadata=None) + assert not attn_metadata.has_spec_dec_saved_state + + +# ---------------- rejection slot buffers + fail-closed guard ---------------- + +R, K, V = 8, 4, 32 # max_num_requests, max_draft_len, vocab_size + + +def _alloc_meta(**over): + base = dict( + use_rejection_sampling=True, + draft_probs=None, + vocab_size=V, + draft_vocab_size=V, + max_num_requests=R, + max_draft_len=K, + num_seq_slots=0, + batch_slot_ids=None, + full_draft_probs=None, + draft_probs_vocab_size=0, + ) + base.update(over) + return types.SimpleNamespace(**base) + + +@_cuda_only +def test_prepare_buffers_allocates_when_enabled(): + """When rejection sampling is enabled, slot buffers are allocated; + full_draft_probs appears only when draft and target vocabs differ.""" + m = _alloc_meta() + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert m.draft_probs is not None and tuple(m.draft_probs.shape) == (R + 1, K, V) + assert m.batch_slot_ids is not None and m.batch_slot_ids.shape[0] == R + assert m.batch_slot_ids.dtype == torch.long + assert m.full_draft_probs is None + assert m.draft_probs_vocab_size == V + + # Distinct draft vocab: full_draft_probs (d2t-expanded) is allocated. + m = _alloc_meta(draft_vocab_size=V - 1) + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert m.full_draft_probs is not None and tuple(m.full_draft_probs.shape) == (R + 1, K, V) + + +@_cuda_only +def test_prepare_buffers_span_seq_slot_pool(): + """Slot-indexed buffers must span the full seq-slot pool (which can exceed + max_num_requests) plus a dummy scratch row placed on the last row.""" + pool = 2 * R + m = _alloc_meta(num_seq_slots=pool) + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert tuple(m.draft_probs.shape) == (pool + 1, K, V) + assert m.dummy_slot_row == pool + + +@_cuda_only +def test_prepare_buffers_noop_when_disabled(): + """With rejection sampling disabled, no buffers are allocated.""" + m = _alloc_meta(use_rejection_sampling=False) + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert m.draft_probs is None + assert m.batch_slot_ids is None + assert m.full_draft_probs is None + + +def _valid_state(): + return types.SimpleNamespace( + draft_probs=torch.empty((R, K, V), device="cuda"), + batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), + is_ragged_verify=False, + ) + + +# num_contexts=0, num_gens=4 +_NUM_CTX = 0 +_BATCH = 4 +_NUM_GENS = _BATCH - _NUM_CTX + + +def _good_args(): + draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") + logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") + return draft_tokens, K, V, _NUM_CTX, _BATCH, logits + + +_DUMMY = object() # _rejection_buffers_valid does not use self + + +def _call(meta, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits): + return SpecWorkerBase._rejection_buffers_valid( + _DUMMY, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits, meta + ) + + +@_cuda_only +def test_guard_true_on_valid_state(): + """The guard accepts a well-formed buffer/argument combination.""" + assert _call(_valid_state(), *_good_args()) is True + + +# Each builder returns (meta, *_call args) for one malformed state that the +# fail-closed guard must reject. +def _case_draft_probs_missing(): + m = _valid_state() + m.draft_probs = None + return (m, *_good_args()) + + +def _case_batch_slot_ids_missing(): + m = _valid_state() + m.batch_slot_ids = None + return (m, *_good_args()) + + +def _case_stored_vocab_nonpositive(): + dt, dl, _sv, nc, bs, lg = _good_args() + return (_valid_state(), dt, dl, 0, nc, bs, lg) + + +def _case_stored_vocab_exceeds_buffer(): + dt, dl, _sv, nc, bs, lg = _good_args() + return (_valid_state(), dt, dl, V + 1, nc, bs, lg) + + +def _case_draft_len_exceeds_buffer(): + # draft_probs has K steps; ask for K+1. + draft_tokens = torch.zeros((_NUM_GENS, K + 1), dtype=torch.int, device="cuda") + logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 2), V), device="cuda") + return (_valid_state(), draft_tokens, K + 1, V, _NUM_CTX, _BATCH, logits) + + +def _case_draft_tokens_wrong_rows(): + dt = torch.zeros((_NUM_GENS + 1, K), dtype=torch.int, device="cuda") + _, dl, sv, nc, bs, lg = _good_args() + return (_valid_state(), dt, dl, sv, nc, bs, lg) + + +def _case_too_few_logits_rows(): + dt, dl, sv, nc, bs, _lg = _good_args() + too_few = torch.zeros((nc + _NUM_GENS, V), device="cuda") # missing +1 each + return (_valid_state(), dt, dl, sv, nc, bs, too_few) + + +def _case_batch_slot_ids_too_short(): + m = _valid_state() + m.batch_slot_ids = torch.arange(_BATCH - 1, device="cuda", dtype=torch.long) + return (m, *_good_args()) + + +@_cuda_only +@pytest.mark.parametrize( + "build_case", + [ + _case_draft_probs_missing, + _case_batch_slot_ids_missing, + _case_stored_vocab_nonpositive, + _case_stored_vocab_exceeds_buffer, + _case_draft_len_exceeds_buffer, + _case_draft_tokens_wrong_rows, + _case_too_few_logits_rows, + _case_batch_slot_ids_too_short, + ], + ids=lambda f: f.__name__.removeprefix("_case_"), +) +def test_guard_false_on_malformed_state(build_case): + """The fail-closed guard rejects every malformed buffer/argument state.""" + assert _call(*build_case()) is False + + +# Fail-closed acceptance dispatch: _accept_draft_tokens() must route to +# strict/base acceptance when the buffers are malformed and to the rejection +# method when the state is valid; the acceptance methods are stubbed to +# record which path ran while the real guard logic is exercised. + + +class _Worker(SpecWorkerBase): + @property + def max_draft_len(self) -> int: + return K + + def _forward_impl(self, *args: object, **kwargs: object) -> None: + raise NotImplementedError + + +def _dispatch_meta(**over): + base = dict( + use_rejection_sampling=True, + is_all_greedy_sample=False, + is_ragged_verify=False, + draft_probs_vocab_size=V, + draft_probs_last_dim=V, + batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), + draft_probs=torch.zeros((R, K, V), device="cuda"), + ) + base.update(over) + return types.SimpleNamespace(**base) + + +def _make_dispatch_worker(): + w = _Worker() + calls = {"base": 0, "rejection": 0} + + def _base(logits, draft_tokens, num_contexts, batch_size, spec_metadata): + calls["base"] += 1 + return ("base", None) + + def _rej(logits, draft_tokens, draft_probs, num_contexts, batch_size, spec_metadata): + calls["rejection"] += 1 + return ("rejection", None) + + w._sample_and_accept_draft_tokens_base = _base + w._sample_and_accept_draft_tokens_rejection = _rej + return w, calls + + +@_cuda_only +@pytest.mark.parametrize( + "meta_over, expected_path", + [ + # Malformed buffers (draft_probs None): _rejection_buffers_valid must + # return False, so acceptance fails closed to base and the rejection + # kernel is skipped. + (dict(draft_probs=None), "base"), + # Valid buffers: the rejection path runs. + (dict(), "rejection"), + # All-greedy batch: rejection is bypassed regardless of buffers. + (dict(is_all_greedy_sample=True), "base"), + ], + ids=[ + "fails_closed_on_malformed_buffers", + "routes_to_rejection_on_valid_state", + "base_when_all_greedy", + ], +) +def test_accept_dispatch(meta_over, expected_path): + """Acceptance dispatch fails closed to base on bad state and routes to + the rejection path only when the buffers are valid and not all-greedy.""" + w, calls = _make_dispatch_worker() + meta = _dispatch_meta(**meta_over) + draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") + logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") + out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) + assert out == (expected_path, None) + other = "rejection" if expected_path == "base" else "base" + assert calls[expected_path] == 1 and calls[other] == 0 + + +# ---------------- group-synchronized is_all_greedy_sample (CPU) ------------- + + +def _fake_request(temperature=None, top_k=None, top_p=None, slot=0): + return types.SimpleNamespace( + sampling_config=types.SimpleNamespace( + temperature=[temperature] if temperature is not None else None, + top_k=[top_k] if top_k is not None else None, + top_p=[top_p] if top_p is not None else None, + ), + state=LlmRequestState.GENERATION_IN_PROGRESS, + py_seq_slot=slot, + ) + + +def _fake_meta(group_all_greedy_sample=None, force_capture=False): + return types.SimpleNamespace( + runtime_draft_len=2, + dummy_slot_row=0, + group_all_greedy_sample=group_all_greedy_sample, + _force_non_greedy_for_capture=force_capture, + ) + + +def _scan(meta, requests): + return SpecMetadata._scan_one_model_sampling(meta, requests) + + +def test_local_value_used_when_no_group_sync(): + """Without a group-synced value, the scan uses the rank-local flag.""" + meta = _fake_meta(group_all_greedy_sample=None) + _scan(meta, [_fake_request(), _fake_request()]) + assert meta.is_all_greedy_sample is True + + _scan(meta, [_fake_request(), _fake_request(temperature=0.8)]) + assert meta.is_all_greedy_sample is False + + +def test_group_override_pulls_greedy_rank_onto_advanced_path(): + """The group AND overrides the rank-local all-greedy flag, and composes + with the warmup capture force without changing the outcome.""" + cases = [ + (False, False, False), + (True, False, True), + (False, True, False), + ] + for group_value, force_capture, expected in cases: + meta = _fake_meta(group_all_greedy_sample=group_value, + force_capture=force_capture) + _scan(meta, [_fake_request(), _fake_request()]) + assert meta.is_all_greedy_sample is expected, (group_value, + force_capture) + + +def test_group_override_survives_rescan(): + """The group override must keep applying on every rescan after the CUDA + graph key is built.""" + meta = _fake_meta(group_all_greedy_sample=False) + for _ in range(3): + _scan(meta, [_fake_request()]) + assert meta.is_all_greedy_sample is False + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py b/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py new file mode 100644 index 000000000000..9387f1732de6 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py @@ -0,0 +1,260 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 + +"""Ragged (per-request) new-token counts for ``compressor_paged_kv_compress``. + +The generation compressor used to take ``NEXT_N`` as a C++ template constant and +derive each request's starting position from it (``sp = kv_len - NEXT_N``). Under +DSpark's ragged verification a request appends only as many tokens as it was +given to verify, so ``next_n`` was demoted to a compile-time *upper bound* and +the true per-request counts now arrive in an optional ``new_tokens_per_seq``. + +The property that matters: a ragged call whose ``new_tokens_per_seq`` is filled +with the uniform value must be **bitwise identical** to passing ``None``. That +is what makes the optional argument safe to add -- it says the ragged code path +did not perturb the uniform one, which every non-DSpark model still uses. + +The second test covers the case the first cannot: genuinely differing counts. +It checks the ragged call against per-request uniform calls, which is the only +independent statement of what "correct" means here. +""" + +import os +from typing import List, Optional + +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("CUDA is required for compressor tests", + allow_module_level=True) + + +def _register_ops() -> None: + """Make ``torch.ops.trtllm.compressor_paged_kv_compress`` available. + + Mirrors ``test_indexer_topk_ragged.py``: importing ``tensorrt_llm`` pulls in + the whole model zoo including optional CuTe-DSL paths that some containers + lack, and this file tests one C++ kernel. + """ + try: + import tensorrt_llm # noqa: F401 + return + except ImportError: + pass + + here = os.path.dirname(os.path.abspath(__file__)) + repo = os.path.abspath(os.path.join(here, *[os.pardir] * 5)) + for candidate in ( + os.path.join(repo, "cpp", "build", "tensorrt_llm", "thop", + "libth_common.so"), + os.path.join(repo, "tensorrt_llm", "libs", "libth_common.so"), + ): + if os.path.exists(candidate): + torch.ops.load_library(candidate) + return + pytest.skip("neither tensorrt_llm nor libth_common.so is importable", + allow_module_level=True) + + +_register_ops() + +HEAD_DIM = 512 +PAGE_SIZE = 64 +COMPRESS_RATIO = 4 +# The paged state is twice as wide as the head for compress_ratio=4: that ratio +# selects the kernel's "overlap" layout, where `STATE_DIM = 2 * HEAD_DIM` (see +# compressorKernels.cu, `IS_OVERLAP`). `kv_score` is twice as wide again -- it +# carries the kv and the score side by side. Getting these wrong does not +# produce a shape error anywhere; the kernel simply reads past the end. +STATE_DIM = 2 * HEAD_DIM if COMPRESS_RATIO == 4 else HEAD_DIM +KV_SCORE_DIM = 2 * STATE_DIM + + +class _CompressorCase: + """Inputs for one ``compressor_paged_kv_compress`` call. + + Built once and reused across the calls being compared, so that a difference + in output can only come from the arguments under test. The paged buffers are + written in place by the kernel, so each call gets a fresh clone. + """ + + def __init__(self, kv_lens: List[int], new_tokens: List[int], + device: torch.device, dtype: torch.dtype, seed: int = 0): + assert len(kv_lens) == len(new_tokens) + self.batch_size = len(kv_lens) + self.kv_lens = list(kv_lens) + self.new_tokens = list(new_tokens) + self.device = device + self.dtype = dtype + + generator = torch.Generator(device="cpu").manual_seed(seed) + + # One block table row per request, sized for the longest sequence. + max_blocks = max( + (kv_len + PAGE_SIZE - 1) // PAGE_SIZE for kv_len in kv_lens) + 1 + num_blocks = self.batch_size * max_blocks + self.block_table = (torch.arange( + num_blocks, dtype=torch.int32).reshape(self.batch_size, + max_blocks).to(device)) + + self.paged_kv = torch.empty([num_blocks, PAGE_SIZE, STATE_DIM], + dtype=dtype).uniform_( + -1, 1, generator=generator).to(device) + self.paged_score = torch.empty([num_blocks, PAGE_SIZE, STATE_DIM], + dtype=dtype).uniform_( + -1, 1, + generator=generator).to(device) + + # kv_score is indexed by the flat token offsets in cu_seq_lens, so it has + # to cover every request's appended tokens. + total_new = sum(new_tokens) + self.kv_score = torch.empty([max(total_new, 1), KV_SCORE_DIM], + dtype=dtype).uniform_( + -1, 1, generator=generator).to(device) + self.ape = torch.empty([COMPRESS_RATIO, STATE_DIM], + dtype=torch.float32).uniform_( + -1, 1, generator=generator).to(device) + + cu_seq_lens = [0] + for count in new_tokens: + cu_seq_lens.append(cu_seq_lens[-1] + count) + self.cu_seq_lens = torch.tensor(cu_seq_lens, + dtype=torch.int32, + device=device) + + # How many compressed tokens each request produces this step: the count + # of compression boundaries its new tokens crossed. + cu_kv_comp = [0] + for kv_len, count in zip(kv_lens, new_tokens): + start = kv_len - count + produced = (kv_len // COMPRESS_RATIO) - (start // COMPRESS_RATIO) + cu_kv_comp.append(cu_kv_comp[-1] + produced) + self.cu_kv_comp = torch.tensor(cu_kv_comp, + dtype=torch.int32, + device=device) + self.total_outputs = cu_kv_comp[-1] + + self.kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32, device=device) + + def run(self, next_n: int, + new_tokens_per_seq: Optional[torch.Tensor]) -> torch.Tensor: + paged_kv = self.paged_kv.clone() + paged_score = self.paged_score.clone() + # Filled, not zeroed: a kernel that skips an output row would otherwise + # be indistinguishable from one that legitimately writes zeros there. + output = torch.full([max(self.total_outputs, 1), HEAD_DIM], + float("nan"), + dtype=self.dtype, + device=self.device) + torch.ops.trtllm.compressor_paged_kv_compress( + self.kv_score, + self.ape, + paged_kv, + paged_score, + self.block_table, + self.block_table, + output, + self.kv_lens_t, + self.cu_seq_lens, + self.cu_kv_comp, + self.batch_size, + PAGE_SIZE, + HEAD_DIM, + COMPRESS_RATIO, + next_n, + new_tokens_per_seq, + ) + return output + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize( + "kv_lens,next_n", + [ + # next_n == 1: the degenerate case the uniform path uses most. + ([256, 257, 512], 1), + # next_n == 6: DeepSeek-V4-Pro-DSpark's full block (max_draft_len=5). + ([256, 300, 512, 1024], 6), + # Lengths deliberately not multiples of compress_ratio, so requests + # produce differing numbers of compressed tokens. + ([255, 258, 261, 1023], 4), + ], +) +def test_uniform_new_tokens_matches_none(kv_lens: List[int], next_n: int, + dtype: torch.dtype): + """Filling ``new_tokens_per_seq`` uniformly must equal passing None. + + This is the differential that licenses the optional argument. If it fails, + the ragged plumbing changed the uniform path -- which every non-DSpark model + on this kernel still relies on. + """ + device = torch.device("cuda") + new_tokens = [next_n] * len(kv_lens) + case = _CompressorCase(kv_lens, new_tokens, device, dtype, seed=1234) + + reference = case.run(next_n, None) + uniform_vector = torch.tensor(new_tokens, dtype=torch.int32, device=device) + ragged = case.run(next_n, uniform_vector) + + assert not torch.isnan(reference).any(), ( + "reference call left output rows unwritten; the test's cu_kv_comp does " + "not match what the kernel produces") + assert torch.equal(reference, ragged), ( + "passing a uniformly-filled new_tokens_per_seq changed the result; the " + "ragged branch is not a strict generalization of the uniform one. Max " + f"abs diff: {(reference.float() - ragged.float()).abs().max().item()}") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_ragged_new_tokens_matches_per_request_uniform(dtype: torch.dtype): + """Genuinely differing counts, checked against single-request uniform calls. + + The test above cannot catch a kernel that ignores ``new_tokens_per_seq`` + entirely -- with a uniform vector that is the correct answer. This one can: + a batch of one request is unambiguously uniform, so running each request + alone with its own ``next_n`` is an independent statement of the expected + result, with no shared code path to the ragged call. + """ + device = torch.device("cuda") + kv_lens = [256, 300, 512, 1024] + new_tokens = [6, 3, 1, 4] + next_n = max(new_tokens) + + batched = _CompressorCase(kv_lens, new_tokens, device, dtype, seed=99) + ragged_out = batched.run(next_n, torch.tensor(new_tokens, + dtype=torch.int32, + device=device)) + assert not torch.isnan(ragged_out).any(), ( + "ragged call left output rows unwritten") + + cu_kv_comp = batched.cu_kv_comp.tolist() + for index, (kv_len, count) in enumerate(zip(kv_lens, new_tokens)): + single = _CompressorCase([kv_len], [count], device, dtype, seed=99) + # Point the one-request case at the same slice of the shared inputs the + # batched call read for this request, so the only difference is the + # batching itself. + token_start = batched.cu_seq_lens[index].item() + single.kv_score = batched.kv_score[token_start:token_start + count] + single.ape = batched.ape + block_row = batched.block_table[index:index + 1] + single.block_table = block_row - block_row.min() + block_start = int(block_row.min().item()) + num_blocks = block_row.shape[1] + single.paged_kv = batched.paged_kv[block_start:block_start + + num_blocks].clone() + single.paged_score = batched.paged_score[block_start:block_start + + num_blocks].clone() + + expected = single.run(count, None) + start, end = cu_kv_comp[index], cu_kv_comp[index + 1] + actual = ragged_out[start:end] + assert torch.equal(expected, actual), ( + f"request {index} (kv_len={kv_len}, new_tokens={count}) differs " + f"between the ragged batch and a uniform batch of one. Max abs " + f"diff: {(expected.float() - actual.float()).abs().max().item()}") diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py new file mode 100644 index 000000000000..9ca99af4af69 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py @@ -0,0 +1,294 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 + +"""Ragged (per-request) query lengths for ``indexer_topk_decode``. + +The uniform kernel reconstructs two things from the scalar ``next_n``: which +request a logits row belongs to (``row // next_n``) and how far back that row +may attend (``seq_len - next_n + row % next_n + 1``). Neither holds once DSpark's +confidence scheduler gives each request its own verify window, so the kernel +takes an optional ``row_kv_lens`` carrying the extent per row. + +The dangerous failure is silent, not loud: a ragged batch whose padded row count +happens to be divisible by the request count satisfies the old +``seq_lens.size(0) * next_n == num_rows`` check, and the kernel then attributes +rows to the wrong requests and masks them at the wrong depth. Every test here is +built to catch that rather than a crash. +""" + +import os + +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("CUDA is required for indexer_topk tests", + allow_module_level=True) + + +def _register_ops() -> None: + """Make ``torch.ops.trtllm.indexer_topk_decode`` available. + + Importing ``tensorrt_llm`` is the usual way, but it drags in the whole model + zoo -- including optional CuTe-DSL code paths that are absent in some + containers. This file tests one C++ kernel, so it falls back to loading the + op library directly rather than being unrunnable whenever an unrelated + optional dependency is missing. + """ + try: + import tensorrt_llm # noqa: F401 + return + except ImportError: + pass + + here = os.path.dirname(os.path.abspath(__file__)) + repo = os.path.abspath(os.path.join(here, *[os.pardir] * 5)) + for candidate in ( + os.path.join(repo, "cpp", "build", "tensorrt_llm", "thop", + "libth_common.so"), + os.path.join(repo, "tensorrt_llm", "libs", "libth_common.so"), + ): + if os.path.exists(candidate): + torch.ops.load_library(candidate) + return + pytest.skip("neither tensorrt_llm nor libth_common.so is importable", + allow_module_level=True) + + +_register_ops() + +INDEX_TOPK = 512 +NUM_COLUMNS = 8192 + +# Every KV length is kept above ``INDEX_TOPK * kMaxBlocksPerRow`` so that each +# block of the split-work path still has at least ``INDEX_TOPK`` candidates to +# rank. Below that a block cannot fill its output slice, the merge pass reads +# whatever was in the scratch buffer, and two runs differ purely by allocator +# luck -- which looks exactly like a kernel bug. +KV_LENS = [2048, 4096, 3000, 2560] + + + +def _assert_same_topk(actual: torch.Tensor, expected: torch.Tensor) -> None: + """Compare top-k *sets* per row, not the raw index arrays. + + The kernel does not promise an output order: the radix path's split-work + merge visits blocks in whatever order they finish, so two identical calls + routinely return the same indices permuted differently. Only the selected + set is part of the contract, and it is what a wrong causal extent would + change. + """ + assert actual.shape == expected.shape, (actual.shape, expected.shape) + actual_sorted = actual.sort(dim=-1).values + expected_sorted = expected.sort(dim=-1).values + torch.testing.assert_close(actual_sorted, expected_sorted, atol=0, rtol=0) + + +def _make_logits(num_rows: int, num_columns: int, + seed: int) -> torch.Tensor: + generator = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(num_rows, + num_columns, + dtype=torch.float32, + device="cuda", + generator=generator) + + +def _row_kv_lens(kv_lens: list[int], verify_lens: list[int]) -> torch.Tensor: + """The extent every row may attend to, matching the kernel's uniform formula. + + Request ``r`` verifying ``v`` positions contributes rows ``o = 0..v-1``, and + row ``o`` sees ``kv_len - v + o + 1`` tokens. Substituting ``v = next_n`` + recovers exactly what the uniform path computes internally, which is what + makes the equivalence test below meaningful. + """ + rows: list[int] = [] + for kv_len, verify_len in zip(kv_lens, verify_lens): + rows.extend(kv_len - verify_len + offset + 1 + for offset in range(verify_len)) + return torch.tensor(rows, dtype=torch.int32, device="cuda") + + +def _run(logits: torch.Tensor, + seq_lens: list[int], + next_n: int, + row_kv_lens: torch.Tensor = None, + compress_ratio: int = 1) -> torch.Tensor: + num_rows = logits.shape[0] + indices = torch.empty(num_rows, + INDEX_TOPK, + dtype=torch.int32, + device="cuda") + # The op refuses to allocate its own split-work scratch (stale pointers are + # a CUDA-graph hazard), so the caller owns it. kMaxBlocksPerRow is small; + # over-allocating is free here and keeps the test independent of the + # dispatcher's blocks-per-row heuristic. + aux_elems = num_rows * 8 * INDEX_TOPK + # Zeroed, not empty: the merge pass reads every slot its blocks could have + # written, so uninitialized scratch makes the result allocator-dependent. + aux_indices = torch.zeros(aux_elems, dtype=torch.int32, device="cuda") + aux_logits = torch.zeros(aux_elems, dtype=torch.float32, device="cuda") + torch.ops.trtllm.indexer_topk_decode( + logits, + torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + indices, + next_n, + INDEX_TOPK, + compress_ratio=compress_ratio, + radix_aux_indices=aux_indices, + radix_aux_logits=aux_logits, + row_kv_lens=row_kv_lens, + ) + return indices + + +@pytest.mark.parametrize("compress_ratio", [1, 4]) +@pytest.mark.parametrize("next_n", [1, 2, 6]) +def test_ragged_with_equal_lens_matches_uniform(next_n, compress_ratio): + """``row_kv_lens`` spelled out by hand must reproduce the uniform result. + + This is the load-bearing test. If the explicit per-row extent and the + kernel's internal ``next_n`` arithmetic ever disagree, every ragged result + is wrong in a way no shape check can see -- so pin them to be bit-identical + on a batch where both are defined. + """ + batch_size = 4 + kv_lens = list(KV_LENS) + logits = _make_logits(batch_size * next_n, NUM_COLUMNS, seed=1234) + + uniform = _run(logits, kv_lens, next_n, compress_ratio=compress_ratio) + ragged = _run(logits, + kv_lens, + next_n=1, + row_kv_lens=_row_kv_lens(kv_lens, [next_n] * batch_size), + compress_ratio=compress_ratio) + + _assert_same_topk(uniform, ragged) + + +def test_genuinely_ragged_matches_per_request_reference(): + """Each request's rows must match running that request on its own. + + A per-request reference is the only check that catches cross-request + misattribution: if row ``i`` were scored against the wrong request's + sequence length, a whole-batch self-consistency check would still pass. + """ + kv_lens = list(KV_LENS) + verify_lens = [6, 3, 1, 4] + total_rows = sum(verify_lens) + logits = _make_logits(total_rows, NUM_COLUMNS, seed=99) + + ragged = _run(logits, + kv_lens, + next_n=1, + row_kv_lens=_row_kv_lens(kv_lens, verify_lens)) + + offset = 0 + for kv_len, verify_len in zip(kv_lens, verify_lens): + rows = logits[offset:offset + verify_len].contiguous() + expected = _run(rows, [kv_len], verify_len) + _assert_same_topk(ragged[offset:offset + verify_len], expected) + offset += verify_len + + +def test_divisible_ragged_batch_is_not_silently_accepted(): + """The shape that used to pass the old check by coincidence. + + ``verify_lens = [6, 4, 1, 1]`` totals 12 rows for 4 requests, so + ``seq_lens.size(0) * next_n == num_rows`` holds for ``next_n == 3`` even + though no request verifies 3 positions. Without ``row_kv_lens`` the kernel + accepts it and computes nonsense; with it, the answer must match the + per-request reference. + """ + kv_lens = list(KV_LENS) + verify_lens = [6, 4, 1, 1] + assert sum(verify_lens) % len(verify_lens) == 0, "test premise" + logits = _make_logits(sum(verify_lens), NUM_COLUMNS, seed=7) + + ragged = _run(logits, + kv_lens, + next_n=1, + row_kv_lens=_row_kv_lens(kv_lens, verify_lens)) + + bogus_next_n = sum(verify_lens) // len(verify_lens) + accidental = _run(logits, kv_lens, bogus_next_n) + assert not torch.equal(ragged.sort(dim=-1).values, + accidental.sort(dim=-1).values), ( + "the ragged result coincided with the misattributed uniform one; the " + "test can no longer detect the bug it exists for") + + +def test_row_kv_lens_length_is_validated(): + """A short ``row_kv_lens`` must raise rather than read out of bounds.""" + logits = _make_logits(8, NUM_COLUMNS, seed=3) + with pytest.raises(RuntimeError, match="one entry per logits row"): + _run(logits, [1024, 2048], next_n=1, + row_kv_lens=torch.tensor([1024] * 4, + dtype=torch.int32, + device="cuda")) + + +def test_cuda_graph_capture_replays_with_new_lens(): + """The ragged path must be capturable and must honor rewritten lengths. + + Capture is the whole point of landing on a token bucket: a ragged step that + cannot replay costs far more than the tokens it trims. Replay with different + contents in the same buffers is what a real step does. + """ + kv_lens = list(KV_LENS) + verify_lens = [6, 3, 1, 4] + total_rows = sum(verify_lens) + logits = _make_logits(total_rows, NUM_COLUMNS, seed=11) + + seq_lens_dev = torch.tensor(kv_lens, dtype=torch.int32, device="cuda") + row_kv_lens_dev = _row_kv_lens(kv_lens, verify_lens) + indices = torch.empty(total_rows, + INDEX_TOPK, + dtype=torch.int32, + device="cuda") + # Capture requires caller-owned scratch: the op refuses to allocate it + # itself precisely because a graph would bake in a stale pointer. + aux_elems = total_rows * 8 * INDEX_TOPK + aux_indices = torch.zeros(aux_elems, dtype=torch.int32, device="cuda") + aux_logits = torch.zeros(aux_elems, dtype=torch.float32, device="cuda") + + def launch(): + torch.ops.trtllm.indexer_topk_decode(logits, + seq_lens_dev, + indices, + 1, + INDEX_TOPK, + compress_ratio=1, + radix_aux_indices=aux_indices, + radix_aux_logits=aux_logits, + row_kv_lens=row_kv_lens_dev) + + # Warm up outside the graph; capture would otherwise record allocator work. + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + launch() + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch() + + # A different split, written in place, must produce the eager answer. + new_verify_lens = [4, 4, 2, 4] + assert sum(new_verify_lens) == total_rows, "replay must keep the shape" + row_kv_lens_dev.copy_(_row_kv_lens(kv_lens, new_verify_lens)) + graph.replay() + torch.cuda.synchronize() + replayed = indices.clone() + + eager = _run(logits, + kv_lens, + next_n=1, + row_kv_lens=_row_kv_lens(kv_lens, new_verify_lens)) + _assert_same_topk(replayed, eager) diff --git a/tests/unittest/evaluate/test_lm_eval_score_selection.py b/tests/unittest/evaluate/test_lm_eval_score_selection.py new file mode 100644 index 000000000000..16ac7da992e6 --- /dev/null +++ b/tests/unittest/evaluate/test_lm_eval_score_selection.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Which entries of an lm-eval result dict count as scores. + +The dict mixes metrics with bookkeeping. ``alias`` holds the task name and +``samples`` holds the example count, neither of which is an accuracy. Treating +them as scores produced two failures in sequence on the same GSM8K run: first a +``TypeError: the resolved dtypes are not compatible with add.reduce`` from +averaging the string, then -- after that was "fixed" by keeping only numeric +entries -- a reported accuracy of 44030.98, because gsm8k's 1319 samples had +been scaled to 131900 and averaged in. The second is the dangerous one: it +cleared a reference of 96.0 and the test passed. + +lm-eval keys metrics as ``","``, so the comma is the structural +discriminator. +""" + +import pytest + +from tensorrt_llm.evaluate.lm_eval import _is_metric_key + + +@pytest.mark.parametrize("key", [ + "exact_match,strict-match", + "exact_match,flexible-extract", + "acc,none", + "acc_norm,none", + "exact_match_stderr,strict-match", +]) +def test_metric_keys_are_recognized(key): + assert _is_metric_key(key) + + +@pytest.mark.parametrize("key", ["alias", "samples"]) +def test_bookkeeping_keys_are_rejected(key): + assert not _is_metric_key(key) + + +def test_gsm8k_shaped_dict_averages_to_the_metrics_only(): + """The exact dict shape that produced 44030.98.""" + scores = { + "alias": "gsm8k", + "exact_match,strict-match": 96.5125, + "exact_match_stderr,strict-match": 0.5053, + "exact_match,flexible-extract": 96.4367, + "exact_match_stderr,flexible-extract": 0.5106, + "samples": 131900, # 1319 examples, already scaled by 100 + } + kept = [ + v for k, v in scores.items() + if _is_metric_key(k) and "_stderr" not in k + ] + assert sorted(kept) == [96.4367, 96.5125] + assert 96.0 < sum(kept) / len(kept) < 97.0 diff --git a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py new file mode 100644 index 000000000000..7e7112e86643 --- /dev/null +++ b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ragged verification must refuse to start without a profiled cost table. + +Without one the planner's budget degenerates to verify-all -- correctly, since +a flat cost model makes every extra verify token look free -- so the ragged +path runs and changes nothing. As a runtime warning that produced 27-minute +runs which "succeeded" having done nothing, detectable only from a counter +(``steps_ragged: 0``) that nobody reads unless they already suspect the +problem. Config errors belong at construction. + +There is no exception: a crafted (non-flat) cost table is how a test makes the +planner trim, which exercises the real confidence-driven path rather than +bypassing it. See +tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py. +""" + +import pytest + +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + +def _cfg(**kwargs): + base = dict(max_draft_len=5, + speculative_model="/nonexistent/model", + enable_confidence_scheduling=True, + enable_ragged_verify=True) + base.update(kwargs) + return DSparkDecodingConfig(**base) + + +def test_ragged_without_a_cost_table_is_rejected(): + with pytest.raises(ValueError, match="requires a profiled cost table"): + _cfg() + # The same gate must open once a table is named -- a gate that rejects + # everything would pass the raise above. + assert _cfg(confidence_sps_table_path="/tmp/table.json").enable_ragged_verify + + +def test_confidence_scheduling_without_ragged_is_rejected(): + """There is no uniform tier ladder to fall back to. + + Scheduling used to have a middle state: pick one verify length for the + whole batch from the captured tiers. It could not act on per-request + confidence -- which is the entire point of a confidence head -- and it + degenerated to the full block whenever acceptance was high, i.e. on every + workload measured on this checkpoint. It is gone, so this combination now + names a path that does not exist and must not silently mean something else. + + The two remaining states are: schedule per request (ragged), or verify the + whole drafted block (confidence scheduling off). + """ + with pytest.raises(ValueError, match="enable_ragged_verify"): + DSparkDecodingConfig(max_draft_len=5, + speculative_model="/nonexistent/model", + enable_confidence_scheduling=True, + enable_ragged_verify=False) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 183481805aed..22fd0058f478 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -3815,3 +3815,63 @@ def test_glm5_resolves_auto_to_python_on_nixl(self): cfg = self._pretrained_config(["GlmMoeDsaForCausalLM"], "glm_moe_dsa") _resolve_transceiver_runtime_auto(args, DeepseekV3ForCausalLM, cfg) assert args.cache_transceiver_config.transceiver_runtime == "PYTHON" + + +class TestDSparkConfidenceScheduling: + """Config surface for DSpark confidence-scheduled verification.""" + + def _cfg(self, **kw): + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + kw.setdefault("max_draft_len", 5) + return DSparkDecodingConfig(**kw) + + def _scheduled(self, **kw): + """Scheduling on, with the companions it now implies. + + Turning scheduling on means ragged -- the uniform tier ladder that + used to sit between "schedule per request" and "verify everything" is + gone -- and ragged means a profiled cost table, without which the + planner's budget degenerates to verify-all. Both are enforced in + llm_args, so a ladder test that sets only the master switch is + testing a configuration that can no longer exist. + """ + kw["enable_confidence_scheduling"] = True + kw.setdefault("enable_ragged_verify", True) + kw.setdefault("confidence_sps_table_path", "/tmp/sps.json") + return self._cfg(**kw) + + def test_defaults_are_off(self): + c = self._cfg() + assert c.enable_confidence_scheduling is False + assert c.confidence_verify_len_tiers is None + # With scheduling off the ladder is just the static length. + assert c.verify_len_tiers == [5] + + def test_knobs_require_the_master_switch(self): + for kw in ({ + "confidence_sts_path": "/tmp/sts.json" + }, { + "confidence_sps_table_path": "/tmp/sps.json" + }, { + "confidence_verify_len_tiers": [1, 5] + }): + with pytest.raises(ValueError, + match="enable_confidence_scheduling"): + self._cfg(**kw) + + def test_default_ladder_is_derived_and_includes_the_full_block(self): + c = self._scheduled() + assert c.verify_len_tiers == [1, 3, 5] + + def test_full_block_is_always_reachable(self): + """The full block is the fallback when confidence is stale; keep it.""" + c = self._scheduled(confidence_verify_len_tiers=[1, 2]) + assert c.verify_len_tiers[-1] == 5 + + def test_ladder_cannot_exceed_max_draft_len(self): + with pytest.raises(ValueError, match="exceeds max_draft_len"): + self._scheduled(confidence_verify_len_tiers=[1, 9]) + + def test_ladder_is_deduped_and_sorted(self): + c = self._scheduled(confidence_verify_len_tiers=[5, 1, 3, 3]) + assert c.verify_len_tiers == [1, 3, 5]