diff --git a/cpp/tensorrt_llm/kernels/IndexerTopK.h b/cpp/tensorrt_llm/kernels/IndexerTopK.h index 656675e641ad..35cb88a41e5d 100644 --- a/cpp/tensorrt_llm/kernels/IndexerTopK.h +++ b/cpp/tensorrt_llm/kernels/IndexerTopK.h @@ -38,10 +38,17 @@ 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 each row's extent is supplied directly instead of reconstructed from `next_n`. +/// Null uses the uniform request-to-row mapping. A negative extent or one whose compressed +/// length exceeds `numColumns` fails closed as an empty row, so device-updated graph inputs +/// cannot read beyond the corresponding logits row. 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 compressRatio = 1, - cudaStream_t const stream = 0); + int const* rowKvLens = nullptr, cudaStream_t const stream = 0); /// bf16 indexer TopK decode — same dispatch tiers as the fp32 entry, except /// the split-work tier is unsupported (the bf16/fp16 entry does not expose @@ -53,12 +60,13 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic /// that regime must use the fp32 entry. 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 = 2048, int const compressRatio = 1, cudaStream_t const stream = 0); + int const next_n, int const topK = 2048, int const compressRatio = 1, 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 compressRatio = 1, cudaStream_t const stream = 0); + int const topK = 2048, 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 d06f4c0385e3..9379dbad7310 100644 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu @@ -298,12 +298,13 @@ __device__ __forceinline__ void decodeSoftmaxVec(void const* __restrict__ paged_ } } -template +template __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, float const* __restrict__ ape, 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 next_n, int page_size, - int max_blocks, int out_elem_bytes) + int max_blocks, 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; @@ -344,6 +345,10 @@ __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]; + if constexpr (USE_RAGGED_TOKENS) + { + next_n = new_tokens_per_seq[batch_idx]; + } int const sp = kv_len - next_n; int const in_off = cu_seq_lens[batch_idx]; int const out_off = cu_kv_comp[batch_idx]; @@ -687,13 +692,17 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo // `const`. cudafe1 rewrites every kernel parameter into a reference when it emits `__wrapper__device_stub_*`, which // turns those otherwise-ignored qualifiers into part of the function type. Clang then rejects the instantiation; // GCC does not diagnose it. -#define INST_DECODE(HD, KV_EB, STATE_EB, CR, NRW) \ - template __global__ void pagedKvCompressKernel(void const* __restrict__, \ +#define INST_DECODE_RAGGED(HD, KV_EB, STATE_EB, CR, NRW, RAGGED) \ + template __global__ void pagedKvCompressKernel(void const* __restrict__, \ float const* __restrict__, void* __restrict__, void* __restrict__, int32_t const* __restrict__, \ int32_t const* __restrict__, void* __restrict__, int32_t const* __restrict__, int32_t const* __restrict__, \ - int32_t const* __restrict__, int, int, int, int); + int32_t const* __restrict__, int, int, int, int, int32_t const* __restrict__); +#define INST_DECODE(HD, KV_EB, STATE_EB, CR, NRW) \ + INST_DECODE_RAGGED(HD, KV_EB, STATE_EB, CR, NRW, false) \ + INST_DECODE_RAGGED(HD, KV_EB, STATE_EB, CR, NRW, true) FOREACH_DECODE_CONFIG(INST_DECODE) #undef INST_DECODE +#undef INST_DECODE_RAGGED // ============================================================================ // Decode Launch Wrapper @@ -711,7 +720,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, - cudaStream_t stream) + int32_t const* new_tokens_per_seq, cudaStream_t stream) { TLLM_CHECK_WITH_INFO( compress_ratio == 4 || compress_ratio == 128, "pagedKvCompressLaunch only supports compress_ratio 4 or 128"); @@ -765,9 +774,18 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k if (head_dim == HD && kv_score_elem_bytes == KV_EB && state_elem_bytes == STATE_EB && compress_ratio == CR \ && num_red_warps == NRW) \ { \ - pagedKvCompressKernel<<>>(kv_score, ape, \ - paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp, \ - next_n, page_size, max_blocks, out_elem_bytes); \ + if (new_tokens_per_seq == nullptr) \ + { \ + pagedKvCompressKernel<<>>( \ + kv_score, ape, paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, \ + cu_kv_comp, next_n, page_size, max_blocks, out_elem_bytes, nullptr); \ + } \ + else \ + { \ + pagedKvCompressKernel<<>>( \ + kv_score, ape, paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, \ + cu_kv_comp, next_n, 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 943ca0205160..3338d65d847e 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 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 b1bce6f2bf67..e796386136a2 100644 --- a/cpp/tensorrt_llm/kernels/indexerTopK.cu +++ b/cpp/tensorrt_llm/kernels/indexerTopK.cu @@ -650,10 +650,11 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowPrefill( } template + typename InputT = float, bool useRaggedKvLens = false> 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) + int* outIndices, int stride0, int stride1, int numColumns, int const topK, int next_n, int compressRatio, + float* outLogits = nullptr, int const numBlocksToMerge = 0, int const* indices = nullptr, + int const* rowKvLens = nullptr) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); @@ -663,9 +664,43 @@ 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; - int rowEnd = actual_kv_len / compressRatio; + int rowEnd = 0; + // The merge kernel consumes only the bounded partial Top-K workspace. It + // must not reinterpret the original request-indexed sequence metadata. + if constexpr (mergeBlocks) + { + rowEnd = numBlocksToMerge * topK; + } + // 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. + else if constexpr (useRaggedKvLens) + { + int const actualKvLen = rowKvLens[rowIdx]; + int const compressedKvLen = actualKvLen / compressRatio; + // rowKvLens is device-produced and may change between graph replays, + // so a host-side value check would synchronize the hot path and bake + // stale values into capture. Validate at the point of use instead. An + // invalid extent becomes an empty row: topKPerRowJob writes -1 (and + // -FLT_MAX for split-work scratch), preventing an out-of-bounds read + // while failing closed under eager execution and graph replay. + bool const validExtent = actualKvLen >= 0 && compressedKvLen <= numColumns; + rowEnd = validExtent ? compressedKvLen : 0; + } + else + { + int seq_len = seqLens[rowIdx / next_n]; + int actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + rowEnd = actual_kv_len / compressRatio; + } // Local pointers to this block if constexpr (!multipleBlocksPerRow && !mergeBlocks) @@ -682,7 +717,6 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I } else if constexpr (mergeBlocks) { - rowEnd = numBlocksToMerge * topK; indices += static_cast(rowIdx) * numBlocksToMerge * topK; outIndices += static_cast(rowIdx) * topK; } @@ -798,7 +832,8 @@ 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 compressRatio, cudaStream_t const stream) + int const stride1, int const next_n, int const topK, int const compressRatio, int const* rowKvLens, + cudaStream_t const stream) { constexpr int kNumThreadsPerBlock = 512; int const blocksPerRow = computeIndexerTopKDecodeBlocksPerRow(numRows, numColumns, splitWorkThreshold); @@ -812,8 +847,11 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic // Single block per row. Below kSortingAlgorithmThreshold use insertion sort, // above use the histogram-radix path. bool const useRadixSort = numColumns >= kSortingAlgorithmThreshold; - auto* kernel_instance = useRadixSort ? &topKPerRowDecode - : &topKPerRowDecode; + auto* kernel_instance = rowKvLens != nullptr + ? (useRadixSort ? &topKPerRowDecode + : &topKPerRowDecode) + : (useRadixSort ? &topKPerRowDecode + : &topKPerRowDecode); cudaLaunchConfig_t config; config.gridDim = numRows; @@ -823,13 +861,15 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic config.numAttrs = 1; config.attrs = attrs; - cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, numColumns, topK, + next_n, compressRatio, nullptr, 0, nullptr, rowKvLens); } else { // Split each row across `blocksPerRow` blocks, then merge with a second pass. - auto* kernel_instance_part1 = &topKPerRowDecode; + auto* kernel_instance_part1 = rowKvLens != nullptr + ? &topKPerRowDecode + : &topKPerRowDecode; cudaLaunchConfig_t config_part1; config_part1.gridDim = dim3(numRows, blocksPerRow); config_part1.blockDim = kNumThreadsPerBlock; @@ -838,8 +878,8 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic config_part1.numAttrs = 1; config_part1.attrs = attrs; - cudaLaunchKernelEx(&config_part1, kernel_instance_part1, logits, seqLens, outIndicesAux, stride0, stride1, topK, - next_n, compressRatio, outLogitsAux, 0, nullptr); + cudaLaunchKernelEx(&config_part1, kernel_instance_part1, logits, seqLens, outIndicesAux, stride0, stride1, + numColumns, topK, next_n, compressRatio, outLogitsAux, 0, nullptr, rowKvLens); constexpr int kNumThreadsPerBlockMerge = 1024; auto* kernel_instance_part2 = &topKPerRowDecode; @@ -851,8 +891,10 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic config_part2.numAttrs = 1; config_part2.attrs = attrs; + // The merge specialization ignores request metadata and reads only the + // bounded partial Top-K workspace. cudaLaunchKernelEx(&config_part2, kernel_instance_part2, outLogitsAux, seqLens, indices, blocksPerRow * topK, 1, - topK, next_n, 1, nullptr, blocksPerRow, outIndicesAux); + blocksPerRow * topK, topK, next_n, 1, nullptr, blocksPerRow, outIndicesAux, nullptr); } sync_check_cuda_error(stream); } @@ -880,7 +922,7 @@ namespace 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 compressRatio, cudaStream_t const stream) + int const compressRatio, int const* rowKvLens, cudaStream_t const stream) { static_assert(std::is_same_v || std::is_same_v, "invokeIndexerTopKDecodeDtype is for bf16/fp16 only"); @@ -891,8 +933,11 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* if (numColumns < kSortingAlgorithmThreshold) { // Insertion sort path — InputT propagated; histogram/sort run on float keys. - auto* kernel_instance = &topKPerRowDecode; + auto* kernel_instance = rowKvLens != nullptr + ? &topKPerRowDecode + : &topKPerRowDecode; cudaLaunchConfig_t config; config.gridDim = numRows; @@ -905,14 +950,17 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* config.numAttrs = 1; config.attrs = attrs; - cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, numColumns, topK, + next_n, compressRatio, nullptr, 0, nullptr, rowKvLens); } else if (numColumns < effectiveSplitWorkThreshold) { // Radix sort path — InputT propagated; histogram/sort run on float keys. - auto* kernel_instance = &topKPerRowDecode; + auto* kernel_instance = rowKvLens != nullptr + ? &topKPerRowDecode + : &topKPerRowDecode; cudaLaunchConfig_t config; config.gridDim = numRows; @@ -925,8 +973,8 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* config.numAttrs = 1; config.attrs = attrs; - cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); + cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, numColumns, topK, + next_n, compressRatio, nullptr, 0, nullptr, rowKvLens); } else { @@ -944,18 +992,18 @@ 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 compressRatio, cudaStream_t const stream) + int const next_n, int const topK, int const compressRatio, int const* rowKvLens, cudaStream_t const stream) { invokeIndexerTopKDecodeDtype<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, - stride0, stride1, next_n, topK, compressRatio, stream); + stride0, stride1, next_n, topK, 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 compressRatio, cudaStream_t const stream) + int const compressRatio, int const* rowKvLens, cudaStream_t const stream) { invokeIndexerTopKDecodeDtype<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, compressRatio, stream); + stride1, next_n, topK, 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 2352f82def77..62e3c2928ff6 100644 --- a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp @@ -37,7 +37,7 @@ 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, int64_t compress_ratio, std::optional const& radix_aux_indices, - std::optional const& radix_aux_logits) + std::optional const& radix_aux_logits, std::optional const& row_kv_lens) { TORCH_CHECK(compress_ratio > 0, "compress_ratio must be greater than 0"); @@ -52,8 +52,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"); @@ -69,6 +78,20 @@ 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(); + } + auto const logits_dtype = logits.scalar_type(); TORCH_CHECK(logits_dtype == at::ScalarType::Float || logits_dtype == at::ScalarType::BFloat16 || logits_dtype == at::ScalarType::Half, @@ -124,21 +147,21 @@ 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), - static_cast(compress_ratio), stream); + static_cast(compress_ratio), rowKvLensPtr, stream); } else if (logits_dtype == at::ScalarType::BFloat16) { tk::invokeIndexerTopKDecode(reinterpret_cast<__nv_bfloat16 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), - 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), static_cast(compress_ratio), - stream); + rowKvLensPtr, stream); } } @@ -188,7 +211,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) { m.def( "indexer_topk_decode(Tensor logits, Tensor seq_lens, Tensor indices, int next_n, int index_topk=2048, " - "int compress_ratio=1, Tensor? radix_aux_indices=None, Tensor? radix_aux_logits=None) -> ()"); + "int compress_ratio=1, 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 5d74d2c8a5d4..bdea0bc07583 100644 --- a/cpp/tensorrt_llm/thop/compressorOp.cpp +++ b/cpp/tensorrt_llm/thop/compressorOp.cpp @@ -37,13 +37,28 @@ 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 = std::numeric_limits::max(); TORCH_CHECK(next_n >= kMinNextN && next_n <= kMaxNextN, "next_n must be in [1, ", kMaxNextN, "] before conversion to int, got ", next_n); + // Ragged batches pass the batch maximum as next_n and the real per-request + // counts here; next_n sizes the launch while these counts bound each request. + 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()); @@ -56,7 +71,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 @@ -147,7 +162,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/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index df7bbf1decf6..b9072a137354 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -334,6 +334,42 @@ state). `TrtllmAttention` path. Its fields are algorithm-specific, not a generic sparse-attention ABI; `VanillaAttention` uses its own per-request contract. +#### 3.1.1 Ragged verification rows + +DSA and DeepSeek-V4 keep request-major metadata for preparation, the indexer, +KV ownership and overlap correction. Python `num_seqs` remains the request +count. For ragged generation only, `presented_token_major()` temporarily +presents one attention row per query token to the MLA RoPE/KV-append call and +to FMHA selection plus dispatch. The native ops derive their sequence count +from those runtime views; each generation row has query length one. + +The presentation swaps `kv_lens_cuda_runtime`, `kv_lens_runtime`, +`prompt_lens_cpu_runtime`, `prompt_lens_cuda_runtime`, +`host_request_types_runtime` and `kv_cache_block_offsets`. It also supplies +the static row-capacity ceiling as `max_num_requests` for workspace sizing. +All views and that capacity are restored on exit, including exceptional exit; +`seq_lens` and request identities are not replaced. + +Preparation initializes both the indexer's per-row causal KV extents and the +attention row extents on device, so prepare-to-forward callers can consume +them immediately. `on_update_kv_lens()` refreshes those same buffers in place +after device-side overlap correction. Device-selected windows must install +their request map and KV corrections before that refresh/replay. Context +attention rows start at zero; generation attention rows start at +`num_contexts`, while generation tokens start at `num_ctx_tokens`. + +The module-level MLA prefix hook matches that presentation too. Ragged +generation uses fixed-address, opt-in prefix buffers over generation rows +(one query token per row), excluding the context prefix. KV prefixes scan +the presented causal device extents, without extra reserved KV tokens. +Prepare/update invalidation rebuilds them after overlap correction; uniform +and context prefix preparation keep their existing paths. + +Uniform indexer scheduling uses the runtime query width, not the persistent +draft-cap allocation width. Ragged generation is not supported with FlashMLA: +its tile schedule uses a request-major scalar query width, so the backend +rejects that combination before scheduling or dispatch. + ### 3.2 KV-cache and decode-time semantics The main question is not just "does the backend read K and V?" but: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py index 5908932ac61d..64d90b780996 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/compressor.py @@ -253,7 +253,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( @@ -272,6 +277,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/backends/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py index 9e3d23886ecc..316897a406a6 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py @@ -41,6 +41,11 @@ class DeepseekV4TrtllmAttentionMetadata(DSAtrtllmAttentionMetadata): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + def _invalidate_mla_scheduler_buffers(self) -> None: + """Invalidate request- and token-prefix buffers derived from sequence lengths.""" + super()._invalidate_mla_scheduler_buffers() + self._fused_q_gen_cu_seqlens_valid = False + def __post_init__(self): super().__post_init__() self.num_total_compressed_tokens = {} @@ -105,6 +110,30 @@ def __post_init__(self): ) self.cu_seq_lens[0] = 0 + self._fused_q_gen_cu_seqlens_valid = False + self.fused_q_gen_cu_seqlens = None + self.gen_new_tokens_per_seq_cuda = None + self.gen_new_tokens_per_seq: Optional[torch.Tensor] = None + if self.enable_ragged_verification: + # The fused Q normalization op consumes request-major token + # prefixes, rather than the sparse-MLA token-major row view. + self.fused_q_gen_cu_seqlens = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences + 1,), + cache_name="fused_q_gen_cu_seqlens", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.fused_q_gen_cu_seqlens[:1].zero_() + # Exact tokens appended by each generation request. + 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, + ) + # 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( @@ -354,6 +383,47 @@ def allocate(name: str, shape: tuple[int, ...], value: int) -> torch.Tensor: read_table_valid=allocate("read_table_valid", (1,), 1), ) + def mla_prepare_fused_q_gen_cu_seqlens(self) -> Optional[torch.Tensor]: + """Return stable request-major query prefixes for ragged fused Q-RoPE. + + The fused Q normalization op accepts token prefixes (not Q rows), so + this deliberately uses request-major ``seq_lens_cuda`` rather than the + one-row-per-token presentation used by sparse MLA. The cumsum writes + into graph-owned storage once per iteration; device-window graph + replay therefore updates values without changing the captured address. + """ + if not self.is_ragged_verify or self.num_generations <= 0: + return None + + num_contexts = self.num_contexts + num_generations = self.num_generations + cu_seqlens = getattr(self, "fused_q_gen_cu_seqlens", None) + seq_lens_cuda = getattr(self, "seq_lens_cuda", None) + if cu_seqlens is None or seq_lens_cuda is None or num_generations + 1 > cu_seqlens.shape[0]: + return None + + if not self._fused_q_gen_cu_seqlens_valid: + torch.cumsum( + seq_lens_cuda[num_contexts : num_contexts + num_generations], + 0, + dtype=torch.int32, + out=cu_seqlens[1 : num_generations + 1], + ) + self._fused_q_gen_cu_seqlens_valid = True + return cu_seqlens[: num_generations + 1] + + def apply_device_ragged_layout( + self, + verify_lens: torch.Tensor, + req_idx: torch.Tensor, + kv_correction: torch.Tensor, + ) -> None: + """Install device-selected windows and invalidate derived Q prefixes.""" + super().apply_device_ragged_layout(verify_lens, req_idx, kv_correction) + # This call precedes the first model layer. Its fixed-address cumsum is + # therefore captured once after the layout update and replayed in order. + self._invalidate_mla_scheduler_buffers() + def prepare_for_indexer_k_cache(self): """Prepare the shared indexer K-cache decode table for DSA kernels.""" # INDEXER_COMPRESS uses shared page indices, so the generic DSA @@ -633,7 +703,8 @@ def validate_sparse_offload_batch(self) -> None: if state is not None: state.is_prefill = True - def prepare(self): + def prepare(self) -> None: + self._validate_ragged_verification() assert self.kv_cache_manager is not None assert self.request_ids is not None self.validate_sparse_offload_batch() @@ -671,6 +742,8 @@ def prepare(self): cached_token_lens = self.cached_token_lens_cpu kv_lens = cached_token_lens[:num_requests] + self.seq_lens_kv[:num_requests] + kv_lens = self._prepare_device_window_kv_lens(kv_lens) + self.cached_token_lens_cuda[:num_requests].copy_( cached_token_lens[:num_requests], non_blocking=True ) @@ -706,10 +779,7 @@ 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 + 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] @@ -764,6 +834,30 @@ def prepare(self): self._compress_ratios_sorted, ) + def _sync_gen_tokens_per_seq(self, num_gen_tokens: int) -> int: + """Synchronize the compressor's scalar and per-request token counts. + + The compressor uses ``num_gen_tokens_per_seq`` as a runtime launch + bound. A ragged batch therefore uses the configured global maximum, + while exact request lengths travel in a device vector. The vector is + copied device-to-device from ``_seq_lens_cuda`` so CUDA-graph replay + observes device-window rewrites instead of reloading the captured host + shape split. + """ + if self.is_ragged_verify and self.num_generations > 0: + self.num_gen_tokens_per_seq = 1 + self.max_draft_tokens + num_contexts = self.num_contexts + num_generations = self.num_generations + gen_new_tokens = self.gen_new_tokens_per_seq_cuda[:num_generations] + gen_new_tokens.copy_(self._seq_lens_cuda[num_contexts : num_contexts + num_generations]) + 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, @@ -839,9 +933,7 @@ 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 - ) + 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/backends/sparse/deepseek_v4/module.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/module.py index b076c4e53a46..b50945c23056 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/module.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/module.py @@ -503,10 +503,11 @@ def _fused_q_rope_specs( """(cos_sin, specs) for the fused Q RoPE, one spec per phase in the batch. Each spec is `(rows, cache_lens, seq_len, cu_q_seqlens)` where `rows` is the - row range of the batch it covers. Context rows take positions from a ragged - token cumsum, generation rows from a uniform query length, so one launch - cannot serve both -- a mixed batch yields two specs and the caller issues one - launch per spec over disjoint row ranges of the same output buffers. + row range of the batch it covers. Context and ragged-generation rows take + positions from separate request-major token cumsums; uniform generation + retains the scalar query length. A mixed batch yields two specs and the + caller issues one launch per spec over disjoint row ranges of the same + output buffers. Returns `(None, [])` when the batch does not qualify. """ @@ -534,16 +535,34 @@ def _fused_q_rope_specs( if num_generations > 0: num_gen_tokens = num_tokens - num_ctx_tokens - if num_gen_tokens <= 0 or num_gen_tokens % num_generations != 0: + if num_gen_tokens <= 0: return None, [] - specs.append( - ( - slice(num_ctx_tokens, num_tokens), - cache_lens[num_contexts:], - num_gen_tokens // num_generations, - None, + if bool(getattr(attn_metadata, "is_ragged_verify", False)): + prep_gen = getattr(attn_metadata, "mla_prepare_fused_q_gen_cu_seqlens", None) + if prep_gen is None: + return None, [] + cu_q_seqlens = prep_gen() + if cu_q_seqlens is None: + return None, [] + specs.append( + ( + slice(num_ctx_tokens, num_tokens), + cache_lens[num_contexts : num_contexts + num_generations], + 0, + cu_q_seqlens, + ) + ) + else: + if num_gen_tokens % num_generations != 0: + return None, [] + specs.append( + ( + slice(num_ctx_tokens, num_tokens), + cache_lens[num_contexts:], + num_gen_tokens // num_generations, + None, + ) ) - ) if not specs: return None, [] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py index f16100b9f3bf..f59fed9faf9f 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/indexer.py @@ -1184,17 +1184,14 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): ) metadata.gen_indexer_kv_lens_cuda_runtime = gen_seq_lens if not metadata.use_expanded_buffers_for_mtp: - next_n_cap = 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_context_lens = metadata.update_indexer_kv_lens_2d(gen_seq_lens) scheduler_metadata_buffer = get_paged_mqa_logits_metadata( gen_seq_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, metadata.num_sms ) metadata.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, non_blocking=True) if metadata.max_draft_tokens > 0: scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata( - metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap], + scheduler_context_lens.contiguous(), _DG_SCHEDULE_BLOCK_KV, metadata.num_sms, ) @@ -1204,7 +1201,7 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): 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) + 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, @@ -1685,11 +1682,12 @@ def sparse_attn_indexer( elif has_decode and not metadata.skip_indexer_for_gen_reqs: # 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: + 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[token_offset : token_offset + num_gen_tokens, ...] @@ -1698,7 +1696,11 @@ 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: + 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 @@ -1726,7 +1728,13 @@ 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 + 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[token_offset : token_offset + num_gen_tokens, ...] # Get k cache and call fp8_paged_mqa_logits / fp8_fp4_paged_mqa_logits @@ -1735,7 +1743,7 @@ def sparse_attn_indexer( k_cache = metadata.kv_cache_manager.get_indexer_k_cache_buffers(self.layer_idx) indexer_max_seq_len = metadata.get_indexer_max_seq_len() - if self.use_cute_dsl_paged_mqa_logits: + 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()` @@ -1890,7 +1898,8 @@ def sparse_attn_indexer( is_prefill=False, sequence_lengths=gen_kv_lens_cuda, scan_lengths=scan_lengths, - next_n=next_n, + row_kv_lens=metadata.ragged_row_kv_lens(num_gen_tokens), + next_n=1 if metadata.is_ragged_verify else next_n, max_seq_len=indexer_max_seq_len, gvr_ext_kwargs=gvr_ext_kwargs, radix_aux_indices=metadata.radix_aux_indices, @@ -1905,9 +1914,18 @@ def sparse_attn_indexer( # Keep the GVR prior current for computed, reused, and dense-skip TopK. if gvr_prior_indices is not None and has_decode: - next_n = num_gen_tokens // num_generations decode_topk = topk_indices_buffer[token_offset : token_offset + num_gen_tokens] - last_mtp_topk = decode_topk[next_n - 1 :: next_n] + if metadata.is_ragged_verify: + last_rows = ( + torch.cumsum(metadata.gen_token_repeats_cuda[:num_generations], dim=0).to( + torch.long + ) + - 1 + ) + last_mtp_topk = decode_topk[last_rows] + else: + next_n = num_gen_tokens // num_generations + last_mtp_topk = decode_topk[next_n - 1 :: next_n] prev_topk_dst = gvr_prior_indices[:num_generations] if do_multi_stream() and self.aux_stream is not None: # Overlap the GVR prior write-back with this layer's sparse attention. diff --git a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py index 6520be118cd4..f4707d530b14 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/dsa/metadata.py @@ -7,7 +7,7 @@ import functools import os from dataclasses import dataclass, field -from typing import TYPE_CHECKING, ClassVar, FrozenSet, List, Optional +from typing import TYPE_CHECKING, ClassVar, FrozenSet, List, Optional, Tuple import torch @@ -68,6 +68,26 @@ def build_req_idx_per_token(seq_lens: torch.Tensor, num_tokens: int) -> torch.Te return torch.searchsorted(cu_seq_lens, token_idx, right=True) +@dataclass(frozen=True) +class TokenMajorGenView: + """One attention row per generation query token. + + Only the MLA RoPE and sparse-MLA generation ops consume this view. The + request-major metadata remains unchanged for every other consumer. + """ + + 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 + # workspace by max_num_requests, so a per-step value would thrash it. + max_num_rows: int + + @dataclass(init=False) class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): """Attention metadata for DSA (Dense Sparse Attention) with indexer state.""" @@ -133,9 +153,100 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): # cross-step state). needs_gvr_prior: bool = field(default=False, init=False) gvr_prior_indices: Optional[torch.Tensor] = field(default=None, init=False) + # Query tokens per generation request this step, in batch order. None + # denotes the uniform runtime stride. + ragged_verify_lens: Optional[List[int]] = None + # Allocate the row-major and token-major buffers used only by ragged + # verification. This is fixed before CUDA-graph metadata is created. + enable_ragged_verification: bool = False + # In device-window mode the host split determines only the captured shape; + # the true windows are installed before replay. + device_windows_mode: bool = False + # Query tokens per generation request for a uniform runtime tier. Zero + # falls back to the static draft-token ceiling. + runtime_tokens_per_gen_step: int = 0 + + @property + def is_ragged_verify(self) -> bool: + return self.ragged_verify_lens is not None + + def _validate_ragged_verification(self) -> None: + """Reject layouts whose stable buffers were not enabled at allocation.""" + if self.ragged_verify_lens is not None and not self.enable_ragged_verification: + raise RuntimeError( + "ragged verification metadata was installed without enabling " + "ragged-verification buffers" + ) + + @property + def gen_token_stride(self) -> int: + stride = int(self.runtime_tokens_per_gen_step) + return stride if stride > 0 else 1 + self.max_draft_tokens + + def update_indexer_kv_lens_2d(self, gen_kv_lens: torch.Tensor) -> torch.Tensor: + """Fill and return the active query-width view of persistent storage.""" + width = min(self.gen_token_stride, self.kv_lens_cuda_2d.shape[1]) + context_lens = self.kv_lens_cuda_2d[: self.num_generations, :width] + context_lens.copy_(gen_kv_lens.unsqueeze(-1).expand(-1, width)) + return context_lens + + def mla_prepare_scheduler_buffers(self, num_heads: int) -> tuple[torch.Tensor, torch.Tensor]: + """Precompute prefixes for the same rows presented to MLA generation.""" + view = self.token_major_gen_view() + if view is None: + return super().mla_prepare_scheduler_buffers(num_heads) + + num_ctx = self.num_contexts + n_gen = view.num_rows - num_ctx + if not self._mla_scheduler_buffers_valid: + # RoPE presents one query token per generation row. The device + # extents are causal, pre-extra-KV lengths, just like the + # cache_seq_lens scanned by the non-precomputed kernel path. + torch.cumsum( + view.prompt_lens_cuda[num_ctx:], + 0, + dtype=torch.int32, + out=self.ragged_mla_cu_q_rows[1 : n_gen + 1], + ) + self.ragged_mla_cu_q_rows[1 : n_gen + 1] *= num_heads + torch.cumsum( + view.sequence_length[num_ctx:], + 0, + dtype=torch.int32, + out=self.ragged_mla_cu_kv_seqlens[1 : n_gen + 1], + ) + self._mla_scheduler_buffers_valid = True + return ( + self.ragged_mla_cu_q_rows[: n_gen + 1], + self.ragged_mla_cu_kv_seqlens[: n_gen + 1], + ) + + def gen_token_repeat_list(self) -> List[int]: + if not self.is_ragged_verify: + return [self.gen_token_stride] * self.num_generations + return list(self.ragged_verify_lens) + + def gen_token_repeats(self, device: Optional[torch.device] = None) -> torch.Tensor: + 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 request values once per generation query token.""" + repeats = self.gen_token_repeat_list() + num_tokens = sum(repeats) + if not self.is_ragged_verify: + return values.repeat_interleave(self.gen_token_stride, dim=dim), num_tokens + 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.""" + self.enable_ragged_verification = bool(kwargs.pop("enable_ragged_verification", False)) sparse_attention_config = kwargs.pop("sparse_attention_config", None) if ( kwargs.get("sparse_metadata_params") is None @@ -177,6 +288,13 @@ def __init__(self, *args, **kwargs): def __post_init__(self): """Allocate indexer K-cache buffers and heuristic TopK metadata.""" super().__post_init__() + if self.is_cross: + # Ragged verification is a target self-attention capability. A + # cross-attention metadata object may be a shallow copy of one, but + # must not inherit its verifier-only buffers or per-step layout. + self.enable_ragged_verification = False + self.ragged_verify_lens = None + self._validate_ragged_verification() if not is_dsa_cache_manager(self.kv_cache_manager): has_deepseek_v4_cache_interface = all( hasattr(self.kv_cache_manager, attr) @@ -240,6 +358,7 @@ def __post_init__(self): self._create_nvfp4_mla_generation_buffers(capture_graph=capture_graph) def prepare(self): + self._validate_ragged_verification() super().prepare() self._invalidate_pool_view_cache() @@ -260,6 +379,8 @@ def prepare(self): else: kv_lens = cached_token_lens + self.seq_lens_kv + kv_lens = self._prepare_device_window_kv_lens(kv_lens) + # For mla_rope_append_paged_kv_assign_q self.prepare_for_mla_rope_append(cached_token_lens, kv_lens) @@ -278,6 +399,15 @@ def prepare(self): # Prepare metadata for indexer Indexer.prepare(metadata=self) + def _prepare_device_window_kv_lens(self, kv_lens: torch.Tensor) -> torch.Tensor: + """Bound host-side KV extents until device-selected windows are installed.""" + if self.device_windows_mode and self.is_ragged_verify: + 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] + return kv_lens + def prepare_for_draft_forward(self) -> dict | None: """Select native DSA indexer metadata for a draft forward.""" # DeepSeek-V4 metadata inherits DSA metadata, but its cache manager uses a @@ -525,6 +655,12 @@ def on_update_kv_lens(self) -> None: self.seq_lens_cuda[: self.num_seqs], self.num_tokens ).to(self.req_idx_per_token.dtype) + # The overlap correction may have changed request KV lengths after + # prepare(); rebuild each ragged row's causal extent on device. + if self.enable_ragged_verification: + 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 and not fused_eligible: seq_lens = self.seq_lens_cuda[: self.num_seqs] # Runtime cached lengths after overlap/spec-dec correction. @@ -579,17 +715,14 @@ def on_update_kv_lens(self) -> None: gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs] gen_indexer_kv_lens = self.get_indexer_kv_lens(gen_kv_lens) self.gen_indexer_kv_lens_cuda_runtime = gen_indexer_kv_lens - next_n_cap = self.kv_lens_cuda_2d.shape[1] - self.kv_lens_cuda_2d[: self.num_generations, :next_n_cap].copy_( - gen_indexer_kv_lens.unsqueeze(-1).expand(-1, next_n_cap) - ) + scheduler_context_lens = self.update_indexer_kv_lens_2d(gen_indexer_kv_lens) scheduler_metadata_buffer = get_paged_mqa_logits_metadata( gen_indexer_kv_lens.view(-1, 1), _DG_SCHEDULE_BLOCK_KV, self.num_sms ) self.scheduler_metadata_buffer.copy_(scheduler_metadata_buffer, non_blocking=True) if self.max_draft_tokens > 0 and not self.use_expanded_buffers_for_mtp: scheduler_metadata_buffer_full_next_n = get_paged_mqa_logits_metadata( - self.kv_lens_cuda_2d[: self.num_generations, :next_n_cap], + scheduler_context_lens.contiguous(), _DG_SCHEDULE_BLOCK_KV, self.num_sms, ) @@ -597,12 +730,8 @@ def on_update_kv_lens(self) -> None: 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() - ) + kv_lens_expanded, num_tokens = self.expand_per_gen_token(gen_indexer_kv_lens) + self.kv_lens_expanded_cuda[:num_tokens].copy_(kv_lens_expanded) scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), _DG_SCHEDULE_BLOCK_KV, @@ -676,7 +805,7 @@ def _run_fused_dsa_decode_metadata(self): def _compute_kv_lens_row_reorder(self) -> None: """Prepare the longest-job-first GVR row order once per forward step.""" - next_n = 1 + self.max_draft_tokens + next_n = self.gen_token_stride if self.needs_gvr_prior and self.num_generations * next_n >= 2 * self.num_sms: gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs] order = torch.argsort(gen_kv_lens, descending=True).to(torch.int32) @@ -1029,16 +1158,24 @@ 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: + """Process-wide draft ceiling used only for persistent buffer sizing.""" + cap = getattr(self, "_draft_alloc_cap", None) + if cap is None: + return self.max_draft_tokens + return max(int(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. - Avoids per-forward .expand().contiguous() allocations that break CUDA - graphs. The buffer is written in-place via .copy_() inside - on_update_kv_lens so its address stays stable across replays. + Preparation and on_update_kv_lens fill the persistent backing storage + in place. Scheduling receives a contiguous view at the active runtime + query width, which may be smaller than the allocation ceiling. """ 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, @@ -1049,7 +1186,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, @@ -1059,10 +1196,34 @@ def create_expanded_buffers(self, capture_graph=False): device="cpu", pin_memory=prefer_pinned(), ) + self.gen_token_repeats_cuda = None + self.host_gen_token_repeats = None + self._ragged_num_rows = 0 + self._attn_num_rows = 0 + self.row_kv_lens_cuda = None + self.row_kv_lens_host = None + self.row_kv_correction_cuda = None + self.row_kv_correction_host = None + self.row_req_idx_cuda = None + self.row_req_idx_host = None + self.attn_row_kv_lens_cuda = None + self.attn_row_kv_correction_cuda = None + self.attn_row_req_idx_cuda = None + self.attn_row_kv_lens_host = None + self.attn_row_kv_correction_host = None + self.attn_row_req_idx_host = None + self.attn_row_request_types_host = None + self.attn_row_prompt_lens_cuda = None + self.attn_row_prompt_lens_cpu = None + self.attn_row_block_offsets = None + self.ragged_mla_cu_q_rows = None + self.ragged_mla_cu_kv_seqlens = None + if self.enable_ragged_verification: + self._create_ragged_verification_buffers(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", @@ -1078,7 +1239,7 @@ def create_expanded_buffers(self, capture_graph=False): self.draft_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.draft_kv_cache_manager.max_blocks_per_seq, ], cache_name="draft_block_table_expanded", @@ -1098,6 +1259,127 @@ def create_expanded_buffers(self, capture_graph=False): capture_graph=capture_graph, ) + def _create_ragged_verification_buffers(self, capture_graph: bool) -> None: + """Allocate persistent row layouts required only by ragged verification.""" + # Stable-address repeat vector used by in-graph ragged expansions. + 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(), + ) + row_cap = self.max_num_sequences * (1 + self._draft_sizing_cap) + for name in ("ragged_mla_cu_q_rows", "ragged_mla_cu_kv_seqlens"): + prefix = self.get_empty( + self.cuda_graph_buffers, + (row_cap + 1,), + cache_name=name, + dtype=torch.int32, + capture_graph=capture_graph, + ) + prefix[:1].zero_() + setattr(self, name, prefix) + self.row_kv_lens_cuda = self.get_empty( + self.cuda_graph_buffers, + (row_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() + ) + self.row_kv_correction_cuda = self.get_empty( + self.cuda_graph_buffers, + (row_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() + ) + self.row_req_idx_cuda = self.get_empty( + self.cuda_graph_buffers, + (row_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() + ) + + # MLA RoPE and sparse-MLA generation consume one row per query token; + # every other consumer keeps using request-major metadata. + attn_rows_cap = self.max_num_sequences * (2 + self._draft_sizing_cap) + 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, + ) + 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() + ) + self.attn_row_request_types_host = torch.ones( + (attn_rows_cap,), dtype=torch.int32, pin_memory=prefer_pinned() + ) + 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() + ) + 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, + ) + # This function is only used to create the expanded buffers when the max_draft_tokens is changed. # TODO: remove this function once fp8_paged_mqa_logits supports an arbitrary number of MTP draft tokens. def update_spec_dec_param( @@ -1130,6 +1412,10 @@ def update_spec_dec_param( rows_per_sequence = max(rows_per_sequence, spec_tree_manager._internal_buf_dim) resize_radix_workspace = rows_per_sequence != self._radix_rows_per_sequence self._radix_rows_per_sequence = rows_per_sequence + self._draft_alloc_cap = max( + int(max_total_draft_tokens), + int(getattr(self, "_draft_alloc_cap", 0) or 0), + ) capture_graph = self.is_cuda_graph max_gen_tokens = self.max_num_sequences * (1 + self.max_draft_tokens) if ( @@ -1137,10 +1423,11 @@ def update_spec_dec_param( and self.nvfp4_mla_fp8_scratch.shape[0] != max_gen_tokens ): self._create_nvfp4_mla_generation_buffers(capture_graph=capture_graph) - if self.kv_lens_cuda_2d.shape[1] != 1 + self.max_draft_tokens: + previous_cap = int(getattr(self, "_expanded_alloc_draft_cap", -1)) + allocation_cap = max(self._draft_sizing_cap, previous_cap) + if allocation_cap > previous_cap: + self._expanded_alloc_draft_cap = allocation_cap 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) if resize_radix_workspace: self._create_radix_aux_buffers(capture_graph) @@ -1390,27 +1677,58 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): # - Drop this sm90 branch (and the expanded buffers) once # fp8_paged_mqa_logits supports an arbitrary next_n on sm90 too. 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 + if not self.is_ragged_verify: + self._ragged_num_rows = 0 + self._attn_num_rows = 0 + else: + # The DSL paged-MQA path reconstructs rows from one scalar window + # and cannot express a per-request split. + use_dsl = False + assert not self.is_spec_decoding_enabled, ( + "ragged verification requires dense spec-decoding metadata to be disabled" + ) + + 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) ) 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: + 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 + ) + 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 ) - self._refresh_expanded_block_table(1 + self.max_draft_tokens) + if self.kv_cache_manager is not None and self.num_generations > 0: + max_len = self.host_indexer_k_cache_block_offsets.shape[1] + gen_blocks = self.host_indexer_k_cache_block_offsets[ + self.num_contexts : self.num_seqs, :max_len + ] + expanded_blocks, _ = self.expand_per_gen_token(gen_blocks, dim=0) + self.host_block_table_expanded[:num_tokens, :max_len].copy_(expanded_blocks) + self.block_table_expanded[:num_tokens].copy_( + self.host_block_table_expanded[:num_tokens], 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 + next_n = self.gen_token_stride kernel_atoms = (1, 2, 3) if self.kv_cache_manager.use_fp4 else (1, 2, 3, 4) gen_kv_lens = self.get_indexer_kv_lens(kv_lens[self.num_contexts : self.num_seqs]) max_ctx = int(gen_kv_lens.max().item()) if gen_kv_lens.numel() else 0 @@ -1433,26 +1751,228 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): self._refresh_expanded_block_table(expand_factor) 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 the causal KV extent of each ragged generation row.""" + 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 {len(gen_kv_lens)}" + ) + 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 has " + f"{expected_gen_tokens} generation tokens" + ) + + rows: List[int] = [] + corrections: List[int] = [] + request_indices: List[int] = [] + for request_idx, (kv_len, verify_len) in enumerate(zip(gen_kv_lens, verify_lens)): + verify_len = int(verify_len) + base = int(kv_len) - verify_len + 1 + rows.extend(range(base, base + verify_len)) + corrections.extend(range(1 - verify_len, 1)) + request_indices.extend([request_idx] * verify_len) + + num_rows = len(rows) + assert num_rows <= self.row_kv_lens_cuda.shape[0] + self._ragged_num_rows = num_rows + self.row_kv_lens_host[:num_rows].copy_(torch.tensor(rows, dtype=torch.int32)) + self.row_kv_correction_host[:num_rows].copy_(torch.tensor(corrections, dtype=torch.int32)) + self.row_req_idx_host[:num_rows].copy_(torch.tensor(request_indices, dtype=torch.long)) + self.row_kv_lens_cuda[:num_rows].copy_(self.row_kv_lens_host[:num_rows], non_blocking=True) + self.row_kv_correction_cuda[:num_rows].copy_( + self.row_kv_correction_host[:num_rows], non_blocking=True + ) + self.row_req_idx_cuda[:num_rows].copy_(self.row_req_idx_host[:num_rows], non_blocking=True) + + def refresh_ragged_row_kv_lens(self) -> None: + """Refresh ragged extents after overlap changes ``kv_lens_cuda``.""" + num_rows = self._ragged_num_rows + if not self.is_ragged_verify or num_rows <= 0: + return + gen_kv_lens = self.kv_lens_cuda[self.num_contexts : self.num_seqs] + row_kv_lens = self.row_kv_lens_cuda[:num_rows] + torch.index_select( + gen_kv_lens, + 0, + self.row_req_idx_cuda[:num_rows], + out=row_kv_lens, + ) + row_kv_lens.add_(self.row_kv_correction_cuda[:num_rows]) + + def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: + """Build the temporary one-row-per-query-token attention view.""" + verify_lens = self.ragged_verify_lens + num_gen_rows = self._ragged_num_rows + if not verify_lens or num_gen_rows <= 0: + self._attn_num_rows = 0 + return + + num_contexts = self.num_contexts + num_rows = num_contexts + num_gen_rows + capacity = self.attn_row_kv_lens_cuda.shape[0] + assert num_rows <= capacity, ( + f"token-major generation needs {num_rows} rows but capacity is {capacity}" + ) + + if num_contexts: + self.attn_row_kv_lens_host[:num_contexts].copy_(kv_lens[:num_contexts].to(torch.int32)) + self.attn_row_kv_correction_host[:num_contexts].zero_() + self.attn_row_req_idx_host[:num_contexts].copy_( + torch.arange(num_contexts, dtype=torch.long) + ) + self.attn_row_prompt_lens_cpu[:num_contexts].copy_( + self.prompt_lens_cpu[:num_contexts].to(torch.int32) + ) + + self.attn_row_kv_lens_host[num_contexts:num_rows].copy_( + self.row_kv_lens_host[:num_gen_rows] + ) + self.attn_row_kv_correction_host[num_contexts:num_rows].copy_( + self.row_kv_correction_host[:num_gen_rows] + ) + torch.add( + self.row_req_idx_host[:num_gen_rows], + num_contexts, + out=self.attn_row_req_idx_host[num_contexts:num_rows], + ) + self.attn_row_prompt_lens_cpu[num_contexts:num_rows].fill_(1) + if num_contexts: + self.attn_row_request_types_host[:num_contexts].zero_() + self.attn_row_request_types_host[num_contexts:num_rows].fill_(1) + + self.attn_row_kv_correction_cuda[:num_rows].copy_( + self.attn_row_kv_correction_host[:num_rows], non_blocking=True + ) + self.attn_row_kv_lens_cuda[:num_rows].copy_( + self.attn_row_kv_lens_host[:num_rows], non_blocking=True + ) + self.attn_row_req_idx_cuda[:num_rows].copy_( + self.attn_row_req_idx_host[:num_rows], non_blocking=True + ) + self.attn_row_prompt_lens_cuda[:num_rows].copy_( + self.attn_row_prompt_lens_cpu[:num_rows], non_blocking=True + ) + self._attn_num_rows = num_rows + self.refresh_token_major_block_table() + + def refresh_token_major_gen_rows(self) -> None: + """Refresh token-major KV extents after overlap correction.""" + num_rows = self._attn_num_rows + if not self.is_ragged_verify or num_rows <= 0: + return + torch.index_select( + self.kv_lens_cuda[: self.num_seqs], + 0, + self.attn_row_req_idx_cuda[:num_rows], + out=self.attn_row_kv_lens_cuda[:num_rows], + ) + self.attn_row_kv_lens_cuda[:num_rows].add_(self.attn_row_kv_correction_cuda[:num_rows]) + + def refresh_token_major_block_table(self) -> None: + """Expand the main attention block table along its sequence axis.""" + num_rows = self._attn_num_rows + if num_rows <= 0 or self.kv_cache_block_offsets is None: + return + source = self.kv_cache_block_offsets + num_pools = min(source.shape[0], self.attn_row_block_offsets.shape[0]) + width = min(source.shape[-1], self.attn_row_block_offsets.shape[-1]) + torch.index_select( + source[:num_pools, :, :, :width], + 1, + self.attn_row_req_idx_cuda[:num_rows], + out=self.attn_row_block_offsets[:num_pools, :num_rows, :, :width], + ) + + def token_major_gen_view(self) -> Optional[TokenMajorGenView]: + """Return the ragged generation view expected by ``trtllm.py``.""" + num_rows = self._attn_num_rows + if not self.is_ragged_verify or num_rows <= 0: + return None + return TokenMajorGenView( + num_rows=num_rows, + sequence_length=self.attn_row_kv_lens_cuda[:num_rows], + host_past_key_value_lengths=self.attn_row_kv_lens_host[:num_rows], + host_context_lengths=self.attn_row_prompt_lens_cpu[:num_rows], + prompt_lens_cuda=self.attn_row_prompt_lens_cuda[:num_rows], + host_request_types=self.attn_row_request_types_host[:num_rows], + kv_cache_block_offsets=self.attn_row_block_offsets[:, :num_rows], + max_num_rows=self.attn_row_kv_lens_cuda.shape[0], + ) + + def ragged_row_kv_lens(self, num_tokens: int) -> Optional[torch.Tensor]: + 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})" + ) + 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: + """Install device-selected windows into stable captured buffers.""" + self._validate_ragged_verification() + if not self.is_ragged_verify: + raise RuntimeError("device ragged layout requires enabled ragged verification") + num_contexts = self.num_contexts + num_rows = self._ragged_num_rows + num_generations = self.num_generations + self.seq_lens_cuda[num_contexts : self.num_seqs] = verify_lens.to(self.seq_lens_cuda.dtype) + self.gen_token_repeats_cuda[:num_generations] = verify_lens.to(torch.int64) + self.row_req_idx_cuda[:num_rows] = req_idx + self.row_kv_correction_cuda[:num_rows] = kv_correction + attn_row_start = num_contexts + attn_row_end = attn_row_start + num_rows + self.attn_row_req_idx_cuda[attn_row_start:attn_row_end] = req_idx + num_contexts + self.attn_row_kv_correction_cuda[attn_row_start:attn_row_end] = kv_correction + token_start = self.num_ctx_tokens + token_end = token_start + num_rows + self.req_idx_per_token[token_start:token_end] = (req_idx + num_contexts).to( + self.req_idx_per_token.dtype + ) + 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[:num_rows, :width] = ( + self.indexer_k_cache_block_offsets[num_contexts : self.num_seqs, :width] + .index_select(0, req_idx) + .clamp_(min=0) + ) def _refresh_expanded_block_table(self, repeat_factor: Optional[int] = None): """Refresh the active cache's expanded INDEX_KEY page table.""" if self.kv_cache_manager is None or self.num_generations == 0: return + use_runtime_layout = False if repeat_factor is None: if self.use_expanded_buffers_for_mtp: - repeat_factor = 1 + self.max_draft_tokens + use_runtime_layout = True elif self.expand_for_dsl and self.dsl_expand_factor > 1: repeat_factor = self.dsl_expand_factor else: return - num_tokens = self.num_generations * repeat_factor 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(repeat_factor, dim=0) + if use_runtime_layout: + expanded_blocks, num_tokens = self.expand_per_gen_token(gen_block_tensor, dim=0) + else: + num_tokens = self.num_generations * repeat_factor + expanded_blocks = gen_block_tensor.repeat_interleave(repeat_factor, dim=0) self.host_block_table_expanded[:num_tokens, :max_len].copy_( expanded_blocks, non_blocking=True ) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 644840aefa53..439f5cc5fd22 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -13,13 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import contextlib import functools import math import os import weakref from dataclasses import dataclass, field -from typing import (TYPE_CHECKING, Any, ClassVar, Dict, FrozenSet, List, - Optional, Tuple) +from typing import (TYPE_CHECKING, Any, ClassVar, Dict, FrozenSet, Iterator, + List, Optional, Tuple) import torch @@ -29,6 +30,7 @@ from ...model_config import ModelConfig from ...speculative.interface import SpecMetadata from ...speculative.spec_tree_manager import SpecTreeManager + from .sparse.dsa.metadata import TokenMajorGenView from tensorrt_llm._utils import get_sm_version, maybe_pin_memory, prefer_pinned from tensorrt_llm.bindings.internal import thop @@ -904,6 +906,69 @@ def restore_after_draft_forward(self, saved_state: dict | None) -> None: """Restore backend state modified for draft-forward execution.""" return None + def token_major_gen_view(self) -> Optional["TokenMajorGenView"]: + """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: Optional["TokenMajorGenView"]) -> Iterator[None]: + """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() # Recomputed on first use this iteration; see mla_prepare_scheduler_buffers. @@ -2072,6 +2137,13 @@ def forward( metadata, TrtllmAttentionMetadata, ) + token_major_view = (metadata.token_major_gen_view() + if forward_args.attention_input_type + == AttentionInputType.generation_only else None) + if metadata.enable_flash_mla and token_major_view is not None: + raise ValueError( + "FlashMLA does not support ragged token-major generation; " + "disable FlashMLA for this layout.") # Cross-attention uses the THOP path; the trtllm-gen backend API does # not carry encoder K/V tensors yet. @@ -2358,22 +2430,29 @@ def forward( assert metadata.kv_cache_manager is None assert metadata.num_contexts == metadata.num_seqs - fmha = self._fmha_manager.select(self, q, k, v, metadata, forward_args) - - if fmha is None: - raise RuntimeError( - "No TRT-LLM attention FMHA library supports this request.") - if metadata.is_cuda_graph or not fmha.supports_workspace_reclamation: - # Conservatively disable reclamation for metadata used by graphs - # or backends that can retain staged workspace state. - metadata.workspace_reclaimable = False - try: - fmha.forward(q, k, v, metadata, forward_args) - except RuntimeError as exc: - log_attention_failure_context( - type(self).__name__, self.layer_idx, metadata, - forward_args.attention_window_size, exc) - raise + # 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. Keep the presentation scoped around both + # library selection and dispatch: support checks read the + # same runtime metadata as the selected FMHA implementation. + with metadata.presented_token_major(token_major_view): + fmha = self._fmha_manager.select(self, q, k, v, metadata, + forward_args) + + if fmha is None: + raise RuntimeError( + "No TRT-LLM attention FMHA library supports this request.") + if metadata.is_cuda_graph or not fmha.supports_workspace_reclamation: + # Conservatively disable reclamation for metadata used by graphs + # or backends that can retain staged workspace state. + metadata.workspace_reclaimable = False + try: + fmha.forward(q, k, v, metadata, forward_args) + except RuntimeError as exc: + log_attention_failure_context( + type(self).__name__, self.layer_idx, metadata, + forward_args.attention_window_size, exc) + raise if self.print_skip_softmax_stat: total_blocks, skipped_blocks = self.skip_softmax_stat @@ -2684,6 +2763,47 @@ def mla_rope_generation( # gates and addresses per token instead of per sequence. helix_tensor_params.append(metadata.helix_local_slots) + # 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, + kv_norm_weight, kv_norm_eps, precomputed_cu_seqlens, + precomputed_fmha_scheduler, kv_only, kv_done_elsewhere, + quant_scale_qkv) + + @contextlib.contextmanager + def presented_token_major_for( + self, metadata: TrtllmAttentionMetadata) -> Iterator[None]: + with metadata.presented_token_major(metadata.token_major_gen_view()): + yield + + def _mla_rope_generation_impl( + self, + metadata: TrtllmAttentionMetadata, + fused_q: Optional[torch.Tensor], + q_pe: Optional[torch.Tensor], + latent_cache: torch.Tensor, + cu_q_seqlens: torch.Tensor, + cu_kv_seqlens: torch.Tensor, + fmha_scheduler_counter: torch.Tensor, + mla_bmm1_scale: Optional[torch.Tensor], + mla_bmm2_scale: Optional[torch.Tensor], + quant_q_buffer: torch.Tensor, + out_scale: Optional[torch.Tensor], + helix_tensor_params: List[Optional[torch.Tensor]], + kv_norm_weight: Optional[torch.Tensor], + kv_norm_eps: float, + precomputed_cu_seqlens: bool, + precomputed_fmha_scheduler: bool, + kv_only: bool, + kv_done_elsewhere: bool, + quant_scale_qkv: Optional[torch.Tensor], + ) -> None: torch.ops.trtllm.mla_rope_generation( fused_q, q_pe, diff --git a/tensorrt_llm/_torch/attention/mla.py b/tensorrt_llm/_torch/attention/mla.py index c6cb160b1276..eb6ad4f97bc9 100644 --- a/tensorrt_llm/_torch/attention/mla.py +++ b/tensorrt_llm/_torch/attention/mla.py @@ -1573,8 +1573,21 @@ 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. Uniform verification keeps the established request-major + # allocation rather than scaling these temporary buffers by K. + cu_seqlens_rows = ( + max(num_seqs, num_tokens) + 1 + if getattr(attn_metadata, "is_ragged_verify", False) + else num_seqs + 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) use_fp8_mla = getattr(self.mqa, "has_fp8_kv_cache", False) or getattr( self.mqa, "has_fp4_kv_cache", False diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index d1b719f3a3fc..b0dd082737ac 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -377,7 +377,8 @@ def _(logits, index_topk, compress_ratio=1, radix_aux_indices=None, - radix_aux_logits=None): + radix_aux_logits=None, + row_kv_lens=None): # In-place operation, no return value (void function) pass diff --git a/tensorrt_llm/_torch/modules/top_k.py b/tensorrt_llm/_torch/modules/top_k.py index dcc100e57091..11fcae60f389 100644 --- a/tensorrt_llm/_torch/modules/top_k.py +++ b/tensorrt_llm/_torch/modules/top_k.py @@ -84,6 +84,7 @@ def forward( row_ends: torch.Tensor | None = None, sequence_lengths: torch.Tensor | None = None, scan_lengths: torch.Tensor | None = None, + row_kv_lens: torch.Tensor | None = None, next_n: int = 1, max_seq_len: int | None = None, gvr_ext_kwargs: dict[str, torch.Tensor | None] | None = None, @@ -100,6 +101,11 @@ def forward( row_ends: Per-row exclusive ends for prefill. sequence_lengths: Per-request logical KV lengths for decode. scan_lengths: Per-request score-column lengths for decode. + row_kv_lens: Optional int32 logical KV lengths in uncompressed + token units, with shape ``[num_rows]`` on ``scores.device``, + for ragged decode. The kernel divides by ``compress_ratio``. + When present, dispatches the CUDA radix kernel; + uniform ``next_n`` implementations cannot represent this map. next_n: Number of decode rows per request. max_seq_len: Maximum decode score width used for GVR kernel tuning. gvr_ext_kwargs: GVR-only keyword arguments. ``gvr_prior_indices`` @@ -127,6 +133,7 @@ def forward( scores, sequence_lengths, scan_lengths, + row_kv_lens, output_indices, next_n, max_seq_len, @@ -244,6 +251,7 @@ def _forward_decode( scores: torch.Tensor, sequence_lengths: torch.Tensor, scan_lengths: torch.Tensor, + row_kv_lens: torch.Tensor | None, output_indices: torch.Tensor, next_n: int, max_seq_len: int | None, @@ -251,6 +259,27 @@ def _forward_decode( radix_aux_indices: torch.Tensor | None, radix_aux_logits: torch.Tensor | None, ) -> torch.Tensor: + if row_kv_lens is not None: + if ( + self.decode_implementation != TopKImplementation.CUDA_RADIX + and not torch.compiler.is_dynamo_compiling() + ): + logger.info_once( + "Ragged per-row KV lengths require CUDA radix decode Top-K; " + f"using it instead of {self.decode_implementation.value}.", + key="ragged_decode_radix_fallback", + ) + return self._forward_decode_radix( + scores, + sequence_lengths, + scan_lengths, + output_indices, + next_n, + radix_aux_indices, + radix_aux_logits, + row_kv_lens=row_kv_lens, + ) + if self.decode_implementation == TopKImplementation.TORCH: return self._forward_decode_torch(scores, scan_lengths, output_indices, next_n) @@ -285,9 +314,12 @@ def _forward_decode_radix( next_n: int, radix_aux_indices: torch.Tensor | None, radix_aux_logits: torch.Tensor | None, + row_kv_lens: torch.Tensor | None = None, ) -> torch.Tensor: - use_cute_dsl = self.decode_implementation == TopKImplementation.CUTE_DSL_RADIX and not ( - self.compress_ratio > 1 and next_n > 1 + use_cute_dsl = ( + row_kv_lens is None + and self.decode_implementation == TopKImplementation.CUTE_DSL_RADIX + and not (self.compress_ratio > 1 and next_n > 1) ) if use_cute_dsl: torch.ops.trtllm.cute_dsl_indexer_topk_decode( @@ -311,6 +343,7 @@ def _forward_decode_radix( compress_ratio=self.compress_ratio, radix_aux_indices=radix_indices, radix_aux_logits=radix_values, + row_kv_lens=row_kv_lens, ) return output_indices diff --git a/tests/unittest/_torch/attention/hw_agnostic/test_ragged_metadata_lifecycle.py b/tests/unittest/_torch/attention/hw_agnostic/test_ragged_metadata_lifecycle.py new file mode 100644 index 000000000000..1027b6b97909 --- /dev/null +++ b/tests/unittest/_torch/attention/hw_agnostic/test_ragged_metadata_lifecycle.py @@ -0,0 +1,355 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU coverage of ragged preparation, update and backend dispatch contracts.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.interface import AttentionInputType +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.module import ( + forward_generation_sparse_attn, +) +from tensorrt_llm._torch.attention.backends.sparse.dsa.indexer import Indexer +from tensorrt_llm._torch.attention.backends.sparse.dsa.metadata import DSAtrtllmAttentionMetadata +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata + +pytestmark = pytest.mark.cpu_only + + +def _metadata(seq_lens: list[int], num_contexts: int = 0) -> DSAtrtllmAttentionMetadata: + metadata = DSAtrtllmAttentionMetadata.__new__(DSAtrtllmAttentionMetadata) + metadata._seq_lens = torch.tensor(seq_lens, dtype=torch.int32) + metadata._seq_lens_cuda = metadata._seq_lens.clone() + metadata._num_contexts = num_contexts + metadata._num_generations = len(seq_lens) - num_contexts + metadata._num_ctx_tokens = sum(seq_lens[:num_contexts]) + metadata._num_tokens = sum(seq_lens) + metadata.max_draft_tokens = 5 + metadata.runtime_tokens_per_gen_step = 0 + metadata._indexer_compress_ratio = 1 + metadata.kv_cache_manager = None + metadata.num_sms = 1 + return metadata + + +@pytest.mark.parametrize("runtime_width", [1, 2, 6]) +def test_schedule_update_keeps_runtime_query_width(monkeypatch, runtime_width: int) -> None: + metadata = _metadata([runtime_width, runtime_width]) + metadata.runtime_tokens_per_gen_step = runtime_width + metadata.kv_lens_cuda = torch.tensor([102, 204], dtype=torch.int32) + metadata.kv_lens_cuda_runtime = metadata.kv_lens_cuda + metadata.kv_lens_cuda_2d = torch.full((2, 6), -1, dtype=torch.int32) + metadata.scheduler_metadata_buffer = torch.empty(1, dtype=torch.int32) + metadata.scheduler_metadata_buffer_full_next_n = torch.empty(1, dtype=torch.int32) + metadata.req_idx_per_token = torch.empty(metadata.num_tokens, dtype=torch.int32) + metadata.gen_kv_indptr = torch.zeros(3, dtype=torch.int64) + metadata.gen_cached_token_indptr = torch.zeros(3, dtype=torch.int64) + metadata.enable_ragged_verification = False + metadata.enable_flash_mla = False + metadata.use_fp8_ds_mla = False + metadata.use_expanded_buffers_for_mtp = False + metadata.expand_for_dsl = False + metadata._invalidate_pool_view_cache = Mock() + metadata._compute_kv_lens_row_reorder = Mock() + metadata.prepare_dense_topk_indices = Mock() + scheduler_inputs = [] + + def schedule(context_lens, block_kv, num_sms): + assert context_lens.is_contiguous() + scheduler_inputs.append(context_lens.clone()) + return torch.zeros(1, dtype=torch.int32) + + monkeypatch.setattr( + "tensorrt_llm._torch.attention.backends.sparse.dsa.indexer.get_paged_mqa_logits_metadata", + schedule, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.get_paged_mqa_logits_metadata", + schedule, + ) + Indexer.prepare_scheduler_metadata(metadata) + metadata.kv_lens_cuda.add_(1) + metadata.on_update_kv_lens() + + assert [tuple(value.shape) for value in scheduler_inputs] == [ + (2, 1), + (2, runtime_width), + (2, 1), + (2, runtime_width), + ] + torch.testing.assert_close( + scheduler_inputs[-1], + metadata.kv_lens_cuda[:, None].expand(-1, runtime_width), + ) + + +@pytest.mark.parametrize("num_contexts", [0, 1]) +def test_prepare_initializes_row_extents_and_refresh_reuses_storage( + monkeypatch, num_contexts: int +) -> None: + metadata = _metadata(([3] if num_contexts else []) + [2, 4], num_contexts) + metadata.ragged_verify_lens = [2, 4] + metadata.enable_ragged_verification = True + metadata.is_spec_decoding_enabled = False + metadata.sparse_metadata_params = SimpleNamespace(use_cute_dsl_paged_mqa_logits=False) + metadata.gen_token_repeats_cuda = torch.empty(2, dtype=torch.int64) + metadata.host_gen_token_repeats = torch.empty(2, dtype=torch.int64) + metadata.kv_lens_expanded_host = torch.empty(6, dtype=torch.int32) + metadata.kv_lens_expanded_cuda = torch.empty(6, dtype=torch.int32) + for name in ("row_kv_lens", "row_kv_correction", "row_req_idx"): + dtype = torch.long if name == "row_req_idx" else torch.int32 + for suffix in ("host", "cuda"): + setattr(metadata, f"{name}_{suffix}", torch.full((6,), -777, dtype=dtype)) + for name in ("attn_row_kv_lens", "attn_row_kv_correction", "attn_row_req_idx"): + dtype = torch.long if name == "attn_row_req_idx" else torch.int32 + for suffix in ("host", "cuda"): + setattr(metadata, f"{name}_{suffix}", torch.full((7,), -777, dtype=dtype)) + metadata.prompt_lens_cpu = torch.tensor(([3] if num_contexts else []) + [100, 200]) + metadata.attn_row_prompt_lens_cpu = torch.empty(7, dtype=torch.int32) + metadata.attn_row_prompt_lens_cuda = torch.empty(7, dtype=torch.int32) + metadata.attn_row_request_types_host = torch.empty(7, dtype=torch.int32) + metadata.attn_row_block_offsets = torch.empty(1, 7, 2, 1, dtype=torch.int32) + metadata.kv_cache_block_offsets = None + monkeypatch.setattr( + "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.get_sm_version", + lambda: 100, + ) + kv_lens = torch.tensor(([3] if num_contexts else []) + [102, 204], dtype=torch.int32) + metadata.prepare_for_spec_decode(kv_lens) + + expected = torch.tensor([101, 102, 201, 202, 203, 204], dtype=torch.int32) + torch.testing.assert_close(metadata.ragged_row_kv_lens(6), expected) + view = metadata.token_major_gen_view() + torch.testing.assert_close(view.sequence_length[num_contexts:], expected) + if num_contexts: + assert view.sequence_length[0] == 3 + row_ptr = metadata.row_kv_lens_cuda.data_ptr() + attention_ptr = metadata.attn_row_kv_lens_cuda.data_ptr() + metadata.kv_lens_cuda = kv_lens.clone() + metadata.kv_lens_cuda[num_contexts:] -= torch.tensor([3, 6]) + metadata.refresh_ragged_row_kv_lens() + metadata.refresh_token_major_gen_rows() + + expected -= torch.tensor([3, 3, 6, 6, 6, 6]) + torch.testing.assert_close(metadata.ragged_row_kv_lens(6), expected) + torch.testing.assert_close(view.sequence_length[num_contexts:], expected) + assert metadata.row_kv_lens_cuda.data_ptr() == row_ptr + assert metadata.attn_row_kv_lens_cuda.data_ptr() == attention_ptr + assert metadata.num_seqs == num_contexts + 2 + + +class _StopBeforeDispatch(Exception): + pass + + +def _prefix_metadata(monkeypatch, num_contexts=0, ragged=True): + seq_lens = ([3] if num_contexts else []) + [1, 2, 4] + metadata = _metadata(seq_lens, num_contexts) + metadata.enable_ragged_verification = ragged + metadata.ragged_verify_lens = [1, 2, 4] if ragged else None + metadata.max_num_sequences = len(seq_lens) + metadata.cuda_graph_buffers = {} + metadata.draft_kv_cache_manager = None + metadata.kv_cache_manager = SimpleNamespace( + num_pools=1, num_attention_op_pools=1, max_blocks_per_seq=1 + ) + allocation_names = [] + + def allocate(buffers, shape, *, cache_name, dtype, capture_graph): + allocation_names.append(cache_name) + return torch.empty(shape, dtype=dtype) + + metadata.get_empty = allocate + monkeypatch.setattr( + "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.prefer_pinned", + lambda: False, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.is_dsa_cache_manager", + lambda value: False, + ) + metadata.create_expanded_buffers() + metadata.kv_cache_manager = None + metadata.mla_cu_q_rows = torch.zeros(len(seq_lens) + 1, dtype=torch.int32) + metadata.mla_cu_kv_seqlens = torch.zeros(len(seq_lens) + 1, dtype=torch.int32) + metadata.mla_ctx_cu_q_seqlens = torch.zeros(len(seq_lens) + 1, dtype=torch.int32) + metadata._mla_ctx_cu_seqlens_valid = False + metadata._mla_scheduler_buffers_valid = False + metadata.prompt_lens_cpu = metadata.seq_lens.clone() + metadata.kv_cache_block_offsets = None + metadata.is_spec_decoding_enabled = False + metadata.sparse_metadata_params = SimpleNamespace(use_cute_dsl_paged_mqa_logits=False) + monkeypatch.setattr( + "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.get_sm_version", + lambda: 100, + ) + kv_lens = torch.tensor(([17] if num_contexts else []) + [101, 202, 304], dtype=torch.int32) + metadata.kv_lens_cuda = kv_lens.clone() + # Reserved extra tokens must not leak into the kernel's causal KV scan. + metadata.kv_lens = kv_lens + 9 + if ragged: + metadata.prepare_for_spec_decode(kv_lens) + return metadata, allocation_names + + +def test_feature_off_keeps_request_major_prefixes_and_allocations(monkeypatch): + metadata, names = _prefix_metadata(monkeypatch, ragged=False) + q_prefix, kv_prefix = metadata.mla_prepare_scheduler_buffers(2) + assert "ragged_mla_cu_q_rows" not in names + assert "ragged_mla_cu_kv_seqlens" not in names + assert q_prefix.data_ptr() == metadata.mla_cu_q_rows.data_ptr() + assert kv_prefix.data_ptr() == metadata.mla_cu_kv_seqlens.data_ptr() + torch.testing.assert_close(q_prefix, torch.tensor([0, 2, 6, 14], dtype=torch.int32)) + torch.testing.assert_close(kv_prefix, torch.tensor([0, 101, 303, 607], dtype=torch.int32)) + + +def test_uniform_step_with_opt_in_keeps_request_major_prefixes(monkeypatch): + metadata, _ = _prefix_metadata(monkeypatch) + metadata.ragged_verify_lens = None + metadata._seq_lens.fill_(2) + metadata._seq_lens_cuda.copy_(metadata._seq_lens) + q_prefix, kv_prefix = metadata.mla_prepare_scheduler_buffers(2) + assert q_prefix.data_ptr() == metadata.mla_cu_q_rows.data_ptr() + torch.testing.assert_close(q_prefix, torch.tensor([0, 4, 8, 12], dtype=torch.int32)) + torch.testing.assert_close(kv_prefix, torch.tensor([0, 101, 303, 607], dtype=torch.int32)) + + +@pytest.mark.parametrize("num_contexts", [0, 1]) +def test_module_precomputes_causal_token_major_prefixes_and_refreshes(monkeypatch, num_contexts): + metadata, names = _prefix_metadata(monkeypatch, num_contexts) + if num_contexts: + torch.testing.assert_close( + metadata.mla_prepare_ctx_cu_seqlens(), torch.tensor([0, 3], dtype=torch.int32) + ) + else: + assert metadata.mla_prepare_ctx_cu_seqlens() is None + observed = [] + + def capture(*values, **controls): + assert controls["precomputed_cu_seqlens"] is True + observed.append( + (values[4].clone(), values[5].clone(), values[4].data_ptr(), values[5].data_ptr()) + ) + raise _StopBeforeDispatch + + monkeypatch.setattr( + "tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.module._get_validated_sm_version", + lambda: 100, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.module._mla_gen_scheduler_scalars", + lambda *unused: (None, None, None), + ) + module = SimpleNamespace( + num_heads_tp=2, + qk_head_dim=4, + qk_nope_head_dim=2, + mqa=SimpleNamespace(has_fp8_kv_cache=False, mla_rope_generation=capture), + kv_cache_dtype="bfloat16", + kv_a_layernorm=SimpleNamespace(variance_epsilon=1e-6), + ) + + def launch(): + with pytest.raises(_StopBeforeDispatch): + forward_generation_sparse_attn( + module, + torch.zeros(7, 8), + torch.zeros(7, 2), + torch.zeros(7, 2), + metadata, + torch.zeros(7, 8), + latent_cache=torch.zeros(7, 4), + ) + + launch() + expected_q = torch.arange(0, 15, 2, dtype=torch.int32) + expected_kv = torch.tensor([101, 201, 202, 301, 302, 303, 304], dtype=torch.int32) + torch.testing.assert_close(observed[-1][0], expected_q) + torch.testing.assert_close(observed[-1][1][1:], expected_kv.cumsum(0, dtype=torch.int32)) + assert observed[-1][1][0] == 0 + assert names.count("ragged_mla_cu_q_rows") == 1 + assert names.count("ragged_mla_cu_kv_seqlens") == 1 + launch() + assert observed[-1][2:] == observed[-2][2:] + + # Exercise the actual overlap hook: refresh device rows and invalidate the + # cached prefix before the next module-level dispatch, without host refresh. + metadata.kv_cache_manager = None + metadata.enable_flash_mla = False + metadata.use_fp8_ds_mla = False + metadata.expand_for_dsl = False + metadata.kv_lens_cuda_2d = torch.empty(metadata.num_generations, 6, dtype=torch.int32) + metadata.scheduler_metadata_buffer = torch.empty(1, dtype=torch.int32) + metadata.scheduler_metadata_buffer_expanded = torch.empty(1, dtype=torch.int32) + metadata.req_idx_per_token = torch.empty(metadata.num_tokens, dtype=torch.int32) + metadata.gen_kv_indptr = torch.zeros(4, dtype=torch.int64) + metadata.gen_cached_token_indptr = torch.zeros(4, dtype=torch.int64) + metadata._invalidate_pool_view_cache = Mock() + metadata._compute_kv_lens_row_reorder = Mock() + metadata.prepare_dense_topk_indices = Mock() + monkeypatch.setattr( + "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.get_paged_mqa_logits_metadata", + lambda *unused: torch.zeros(1, dtype=torch.int32), + ) + metadata.kv_lens_cuda[num_contexts:] -= torch.tensor([2, 3, 6]) + metadata.on_update_kv_lens() + assert metadata._mla_scheduler_buffers_valid is False + launch() + expected_kv -= torch.tensor([2, 3, 3, 6, 6, 6, 6]) + torch.testing.assert_close(observed[-1][0], expected_q) + torch.testing.assert_close(observed[-1][1][1:], expected_kv.cumsum(0, dtype=torch.int32)) + assert observed[-1][2:] == observed[-2][2:] + assert metadata.num_seqs == num_contexts + 3 + + +@pytest.mark.parametrize("metadata_valid", [False, True]) +def test_flash_mla_rejects_token_major_generation_before_dispatch(metadata_valid: bool) -> None: + metadata = TrtllmAttentionMetadata.__new__(TrtllmAttentionMetadata) + metadata.enable_flash_mla = True + metadata._flash_mla_metadata_valid = metadata_valid + metadata.token_major_gen_view = Mock(return_value=object()) + attention = SimpleNamespace( + sparse_params=None, + is_mla_enable=False, + create_output=Mock(side_effect=_StopBeforeDispatch), + ) + metadata.use_paged_context_fmha = False + + with pytest.raises(ValueError, match="FlashMLA.*ragged"): + TrtllmAttention.forward( + attention, + torch.empty(3, 8), + None, + None, + metadata, + attention_input_type=AttentionInputType.generation_only, + ) + attention.create_output.assert_not_called() + + +@pytest.mark.parametrize("flash_mla,token_major", [(True, False), (False, True)]) +def test_uniform_flash_mla_and_ragged_native_pass_the_guard( + flash_mla: bool, token_major: bool +) -> None: + metadata = TrtllmAttentionMetadata.__new__(TrtllmAttentionMetadata) + metadata.enable_flash_mla = flash_mla + metadata.token_major_gen_view = Mock(return_value=object() if token_major else None) + metadata.use_paged_context_fmha = False + attention = SimpleNamespace( + sparse_params=None, + is_mla_enable=False, + create_output=Mock(side_effect=_StopBeforeDispatch), + ) + with pytest.raises(_StopBeforeDispatch): + TrtllmAttention.forward( + attention, + torch.empty(3, 8), + None, + None, + metadata, + attention_input_type=AttentionInputType.generation_only, + ) + attention.create_output.assert_called_once() diff --git a/tests/unittest/_torch/attention/hw_agnostic/test_trtllm_mla_rope_forwarding.py b/tests/unittest/_torch/attention/hw_agnostic/test_trtllm_mla_rope_forwarding.py new file mode 100644 index 000000000000..026da392e8fd --- /dev/null +++ b/tests/unittest/_torch/attention/hw_agnostic/test_trtllm_mla_rope_forwarding.py @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + +@pytest.mark.parametrize("helix_spec_tokens_valid", [False, True]) +def test_mla_rope_generation_forwards_kernel_controls(helix_spec_tokens_valid: bool) -> None: + metadata = SimpleNamespace( + kv_cache_manager=object(), + max_seq_len=1024, + helix_position_offsets=object(), + helix_is_inactive_rank=object(), + _helix_spec_tokens_valid=helix_spec_tokens_valid, + helix_local_slots=object(), + ) + dispatch = Mock() + attention = SimpleNamespace( + is_mla_enable=True, + mla_params=object(), + uses_fp4_mla_attention=False, + _ensure_rope_table_size=Mock(), + presented_token_major_for=Mock(return_value=nullcontext()), + _mla_rope_generation_impl=dispatch, + ) + tensors = [object() for _ in range(9)] + weight, scale = object(), object() + controls = dict( + kv_norm_weight=weight, + kv_norm_eps=1e-5, + precomputed_cu_seqlens=True, + precomputed_fmha_scheduler=True, + kv_only=True, + kv_done_elsewhere=False, + quant_scale_qkv=scale, + ) + TrtllmAttention.mla_rope_generation( + attention, + *tensors[:3], + metadata, + *tensors[3:], + **controls, + ) + attention._ensure_rope_table_size.assert_called_once_with(1024) + attention.presented_token_major_for.assert_called_once_with(metadata) + dispatch.assert_called_once() + assert dispatch.call_args.args[0] is metadata + assert dispatch.call_args.args[1:10] == tuple(tensors) + helix_tensor_params = [metadata.helix_position_offsets, metadata.helix_is_inactive_rank] + if helix_spec_tokens_valid: + helix_tensor_params.append(metadata.helix_local_slots) + assert dispatch.call_args.args[11] == helix_tensor_params + assert dispatch.call_args.args[-7:] == tuple(controls.values()) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py index 8b4c05a84486..bbf597aa37e4 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py @@ -130,6 +130,7 @@ def __init__( self.max_ctx_compressed_tokens = max_ctx_compressed_tokens self.compressed_mask_cuda = compressed_mask_cuda self.num_gen_tokens_per_seq = 0 # Set by caller + self.gen_new_tokens_per_seq = None self.kv_lens_cuda_runtime = None # Set by caller self.cached_token_lens_cuda = None # Set by caller diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py index b712aad2e4a8..7998f8750142 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_offload.py @@ -17,7 +17,7 @@ import math from dataclasses import replace -from types import SimpleNamespace +from types import MethodType, SimpleNamespace from unittest.mock import Mock, patch import pytest @@ -356,13 +356,15 @@ def _host_batch(*, prefill=False): manager.compute_sliding_block_tables = Mock() for cache in manager.kv_cache_map.values(): cache.history_length = 0 - return SimpleNamespace( + metadata = SimpleNamespace( sparse_metadata_params=SimpleNamespace(enable_kv_cache_offload=True), sparse_offload_state=SimpleNamespace(prepared=True, is_prefill=False), kv_cache_manager=manager, request_ids=[20, 10], beam_width=1, max_draft_tokens=0, + enable_ragged_verification=False, + ragged_verify_lens=None, draft_kv_cache_manager=None, num_contexts=2 if prefill else 0, num_generations=0 if prefill else 2, @@ -371,6 +373,10 @@ def _host_batch(*, prefill=False): prompt_lens=[8, 12], kv_cache_params=SimpleNamespace(num_cached_tokens_per_seq=[0, 0]), ) + metadata._validate_ragged_verification = MethodType( + DeepseekV4TrtllmAttentionMetadata._validate_ragged_verification, metadata + ) + return metadata @pytest.mark.cpu_only 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 cf9af1e5e402..b04787ed8944 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.backends.interface import ( - AttentionForwardArgs, AttentionInputType, MLAParams, PositionalEmbeddingParams, @@ -39,7 +37,6 @@ DeepseekV4TrtllmAttention, DeepseekV4TrtllmAttentionMetadata, ) -from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.cache_manager import get_token_bytes from tensorrt_llm._torch.attention.backends.sparse.params import SparseBackendForwardArgs from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest @@ -105,6 +102,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): @@ -398,6 +400,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, @@ -408,23 +437,27 @@ 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( - -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) + fused_q_seq = fused_q_rot[lo:hi].unflatten(-1, [num_heads, fused_head_dim]) + 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) @@ -464,27 +497,29 @@ 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( - -1, [num_heads, fused_head_dim] - ) + 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( - -2 - ) # [seq_len_q, 1, head_dim] + new_token_latent = new_latent_cache[lo:hi].unsqueeze(-2) # [seq_len_q, 1, head_dim] # Past latent cache past_latent = latent_cache_ref[total_past_tokens : total_past_tokens + past_len].unsqueeze( @@ -509,7 +544,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()] @@ -546,11 +581,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: @@ -619,7 +664,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 @@ -670,8 +715,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), @@ -702,75 +746,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(), @@ -817,10 +793,45 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): @skip_pre_blackwell @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`. Cover every context layout at width one, then the +# lower and upper speculative widths on the heterogeneous page-boundary layout. +# Pinning every case at width one left width-dependent strides 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 original Cartesian 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( + "context_lengths,generation_seq_len_q", + [ + ([4399], 1), + ([14, 508, 3947], 1), + ([2, 1406, 3327], 1), + ([14, 508, 3947], 2), + ([14, 508, 3947], 6), + ([14, 508, 3947], [1, 2, 4]), + ([14, 508, 3947], [2, 2, 2]), + ], +) +def test_deepseek_v4_sparse_mla(context_lengths: List[int], generation_seq_len_q: int | List[int]): + num_generation_steps = 2 + is_ragged = isinstance(generation_seq_len_q, list) + generation_lens = ( + generation_seq_len_q + if isinstance(generation_seq_len_q, list) + else [generation_seq_len_q] * len(context_lengths) + ) + assert len(generation_lens) == len(context_lengths) + max_generation_width = max(generation_lens) scenario = Scenario() device = torch.device("cuda") dtype = scenario.dtype @@ -838,7 +849,7 @@ def test_deepseek_v4_sparse_mla(context_lengths: List[int], num_generation_steps num_heads = 64 if not rope_append else scenario.num_heads batch_size = len(context_lengths) max_context_len = max(context_lengths) - max_seq_len = max_context_len + (num_generation_steps + 1) * generation_seq_len_q + max_seq_len = max_context_len + (num_generation_steps + 1) * max_generation_width total_ctx_tokens = sum(context_lengths) torch.manual_seed(42) @@ -957,11 +968,11 @@ def yarn_get_mscale(scale=1, mscale=1): torch.cat( [ torch.empty( - [generation_seq_len_q, kv_lora_rank], + [query_len, kv_lora_rank], dtype=dtype, device=device, ).uniform_(-1, 1) - for _ in context_lengths + for query_len in generation_lens ] ) for _ in range(num_generation_steps) @@ -970,11 +981,11 @@ def yarn_get_mscale(scale=1, mscale=1): torch.cat( [ torch.empty( - [generation_seq_len_q, qk_rope_head_dim], + [query_len, qk_rope_head_dim], dtype=dtype, device=device, ).uniform_(-1, 1) - for _ in context_lengths + for query_len in generation_lens ] ) for _ in range(num_generation_steps) @@ -983,11 +994,11 @@ def yarn_get_mscale(scale=1, mscale=1): torch.cat( [ torch.empty( - [generation_seq_len_q, num_heads, kv_lora_rank], + [query_len, num_heads, kv_lora_rank], dtype=dtype, device=device, ).uniform_(-1, 1) - for _ in context_lengths + for query_len in generation_lens ] ) for _ in range(num_generation_steps) @@ -996,11 +1007,11 @@ def yarn_get_mscale(scale=1, mscale=1): torch.cat( [ torch.empty( - [generation_seq_len_q, num_heads, qk_rope_head_dim], + [query_len, num_heads, qk_rope_head_dim], dtype=dtype, device=device, ).uniform_(-1, 1) - for _ in context_lengths + for query_len in generation_lens ] ) for _ in range(num_generation_steps) @@ -1144,10 +1155,14 @@ 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) + for request, query_len in zip(requests, generation_lens): + _allocate_kv_cache_for_generation(cache_manager, [request], num_new_tokens=query_len) - 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] + cached_lens = [ + ctx_len + step * query_len + for ctx_len, query_len in zip(context_lengths, generation_lens) + ] + kv_lens = [cl + query_len for cl, query_len in zip(cached_lens, generation_lens)] # Grow compressed ref data: as KV length increases, new compressed # tokens appear. Simulate this by appending random data. @@ -1164,25 +1179,67 @@ def yarn_get_mscale(scale=1, mscale=1): compress_ref_data[layer_idx], ) - gen_seq_lens = torch.tensor([generation_seq_len_q] * batch_size, dtype=torch.int) - total_gen_tokens = batch_size * generation_seq_len_q - gen_metadata = DeepseekV4TrtllmAttentionMetadata( - seq_lens=gen_seq_lens, - request_ids=request_ids, - max_num_requests=batch_size, - num_contexts=0, - prompt_lens=context_lengths, - max_num_tokens=total_gen_tokens, - kv_cache_manager=cache_manager, - kv_cache_params=KVCacheParams( - use_cache=True, - num_cached_tokens_per_seq=cached_lens, - ), - mapping=mapping, - enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), - sparse_attention_config=sparse_config, + gen_seq_lens = torch.tensor(generation_lens, dtype=torch.int) + total_gen_tokens = sum(generation_lens) + + def create_gen_metadata(ragged: bool) -> DeepseekV4TrtllmAttentionMetadata: + metadata = DeepseekV4TrtllmAttentionMetadata( + seq_lens=gen_seq_lens, + request_ids=request_ids, + max_num_requests=batch_size, + num_contexts=0, + prompt_lens=context_lengths, + max_num_tokens=batch_size * MAX_GENERATION_SEQ_LEN_Q, + kv_cache_manager=cache_manager, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=cached_lens, + ), + mapping=mapping, + enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), + sparse_attention_config=sparse_config, + enable_ragged_verification=ragged, + ) + # Keep the persistent allocation ceiling independent of this + # step's query lengths, as in update_spec_dec_param. + metadata.max_draft_tokens = MAX_GENERATION_SEQ_LEN_Q - 1 + metadata.runtime_tokens_per_gen_step = max_generation_width + metadata.ragged_verify_lens = generation_lens.copy() if ragged else None + metadata._create_kv_lens_2d_buffer() + metadata.create_expanded_buffers() + for ratio, position_ids in metadata.compressed_position_ids_cuda.items(): + required_capacity = batch_size * math.ceil(MAX_GENERATION_SEQ_LEN_Q / ratio) + assert position_ids.numel() >= required_capacity + metadata.prepare() + return metadata + + gen_metadata = create_gen_metadata(is_ragged) + if is_ragged: + expected_rows = torch.tensor( + [ + cl + j + 1 + for cl, length in zip(cached_lens, generation_lens) + for j in range(length) + ], + dtype=torch.int32, + device=device, + ) + torch.testing.assert_close( + gen_metadata.ragged_row_kv_lens(total_gen_tokens), expected_rows + ) + torch.testing.assert_close( + gen_metadata.token_major_gen_view().sequence_length, expected_rows + ) + gen_metadata.on_update_kv_lens() + torch.testing.assert_close( + gen_metadata.ragged_row_kv_lens(total_gen_tokens), expected_rows + ) + torch.testing.assert_close( + gen_metadata.token_major_gen_view().sequence_length, expected_rows + ) + uniform_metadata = ( + create_gen_metadata(False) if is_ragged and len(set(generation_lens)) == 1 else None ) - gen_metadata.prepare() for layer_idx in TEST_LAYERS: ratio = scenario.compress_ratios[layer_idx] @@ -1195,10 +1252,20 @@ def yarn_get_mscale(scale=1, mscale=1): k_pe = inp["gen_k_pe_list"][step] latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) + if uniform_metadata is not None: + uniform_q = fused_q.clone() + uniform_q_pe = q_pe.clone() + uniform_latent = latent_cache.clone() + # Prepare generation-specific tensors - num_seqs = gen_metadata.kv_lens_cuda_runtime.size(0) - cu_q_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=device) - cu_kv_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=device) + # New ragged cases use the metadata prefix hook consumed by the + # module; retain the existing in-kernel preparation cases too. + if is_ragged: + cu_q_seqlens, cu_kv_seqlens = gen_metadata.mla_prepare_scheduler_buffers(num_heads) + else: + num_seqs = gen_metadata.kv_lens_cuda_runtime.size(0) + cu_q_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=device) + cu_kv_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=device) fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=device) layers[layer_idx].mla_rope_generation( @@ -1212,13 +1279,22 @@ def yarn_get_mscale(scale=1, mscale=1): None, # mla_bmm1_scale None, # mla_bmm2_scale None, # quant_q_buffer + precomputed_cu_seqlens=is_ragged, ) - # 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, query_len in zip(cached_lens, generation_lens) + for j in range(query_len) + ], ratio, scenario.index_topk, device, @@ -1238,6 +1314,37 @@ def yarn_get_mscale(scale=1, mscale=1): sparse_backend_args=SparseBackendForwardArgs(topk_indices=topk_indices), ) + if uniform_metadata is not None: + uniform_cu_q = torch.empty(batch_size + 1, dtype=torch.int32, device=device) + uniform_cu_kv = torch.empty_like(uniform_cu_q) + uniform_counter = torch.empty_like(fmha_scheduler_counter) + layers[layer_idx].mla_rope_generation( + uniform_q, + uniform_q_pe, + uniform_latent, + uniform_metadata, + uniform_cu_q, + uniform_cu_kv, + uniform_counter, + None, + None, + None, + ) + uniform_result = layers[layer_idx].forward( + uniform_q, + None, + None, + uniform_metadata, + attention_input_type=AttentionInputType.generation_only, + latent_cache=uniform_latent, + q_pe=uniform_q_pe, + cu_q_seqlens=uniform_cu_q, + cu_kv_seqlens=uniform_cu_kv, + fmha_scheduler_counter=uniform_counter, + sparse_backend_args=SparseBackendForwardArgs(topk_indices=topk_indices), + ) + torch.testing.assert_close(result, uniform_result, atol=0.2, rtol=0.02) + # Reference: apply RoPE separately, then compute attention fused_q_rot, new_latent = _rotate_gen_inputs( fused_q, @@ -1249,6 +1356,7 @@ def yarn_get_mscale(scale=1, mscale=1): num_heads, kv_lora_rank, qk_rope_head_dim, + seq_lens_q=generation_lens, ) ref_result, new_latent_cache = calculate_deepseek_v4_ref_gen_sparse( fused_q_rot, @@ -1266,6 +1374,7 @@ def yarn_get_mscale(scale=1, mscale=1): q_scaling, ratio, attn_sink=attn_sinks[layer_idx], + seq_lens_q=generation_lens, ) latent_cache_ref_all[layer_idx] = new_latent_cache diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_ragged_fused_q_rope_specs.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_ragged_fused_q_rope_specs.py new file mode 100644 index 000000000000..d44330524748 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_ragged_fused_q_rope_specs.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from types import MethodType, SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import ( + DeepseekV4TrtllmAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.module import _fused_q_rope_specs + + +class _Mqa: + def __init__(self): + self.rotary_cos_sin = object() + + def _ensure_rope_table_size(self, size): + pass + + +def _ragged_metadata(verify_lens, *, context_lens=(), kv_lens=None): + num_generations = len(verify_lens) + if kv_lens is None: + kv_lens = [1000 + index + length for index, length in enumerate(verify_lens)] + metadata = SimpleNamespace( + kv_lens_cuda_runtime=torch.tensor(kv_lens, dtype=torch.int32), + num_ctx_tokens=sum(context_lens), + num_tokens=sum(context_lens) + sum(verify_lens), + max_seq_len=max(kv_lens), + is_ragged_verify=True, + num_contexts=len(context_lens), + num_generations=num_generations, + seq_lens_cuda=torch.tensor([*context_lens, *verify_lens], dtype=torch.int32), + fused_q_gen_cu_seqlens=torch.zeros(num_generations + 1, dtype=torch.int32), + _fused_q_gen_cu_seqlens_valid=False, + ) + metadata.mla_prepare_fused_q_gen_cu_seqlens = MethodType( + DeepseekV4TrtllmAttentionMetadata.mla_prepare_fused_q_gen_cu_seqlens, metadata + ) + return metadata + + +@pytest.mark.parametrize( + "verify_lens", + ([3] * 64 + [5] * 64, [4] * 64 + [6] * 64, [5] * 64 + [6] * 64), +) +def test_heterogeneous_g128_generation_contract(verify_lens): + mla = SimpleNamespace(mqa=_Mqa()) + metadata = _ragged_metadata(verify_lens) + cos_sin, specs = _fused_q_rope_specs(mla, metadata, 0, len(verify_lens)) + + assert cos_sin is mla.mqa.rotary_cos_sin + assert len(specs) == 1 + rows, cache_lens, seq_len, cu_q_seqlens = specs[0] + assert (rows.start, rows.stop) == (0, sum(verify_lens)) + torch.testing.assert_close(cache_lens, metadata.kv_lens_cuda_runtime) + assert seq_len == 0 + assert cu_q_seqlens.tolist() == [0, *torch.tensor(verify_lens).cumsum(0).tolist()] + + +def test_uniform_k5_keeps_scalar_generation_spec(): + metadata = SimpleNamespace( + kv_lens_cuda_runtime=torch.full((128,), 1006, dtype=torch.int32), + num_ctx_tokens=0, + num_tokens=768, + max_seq_len=1006, + is_ragged_verify=False, + ) + _, specs = _fused_q_rope_specs(SimpleNamespace(mqa=_Mqa()), metadata, 0, 128) + assert len(specs) == 1 + rows, cache_lens, seq_len, cu_q_seqlens = specs[0] + assert (rows.start, rows.stop) == (0, 768) + assert cache_lens.tolist() == [1006] * 128 + assert seq_len == 6 + assert cu_q_seqlens is None + + +def test_mixed_batch_uses_disjoint_context_and_generation_prefixes(): + ctx_prefix = torch.tensor([0, 2, 5], dtype=torch.int32) + metadata = _ragged_metadata([2, 1, 3], context_lens=[2, 3], kv_lens=[102, 203, 302, 401, 503]) + metadata.mla_prepare_ctx_cu_seqlens = lambda: ctx_prefix + _, specs = _fused_q_rope_specs(SimpleNamespace(mqa=_Mqa()), metadata, 2, 3) + assert len(specs) == 2 + ctx_rows, ctx_cache, ctx_seq_len, ctx_cu_q = specs[0] + gen_rows, gen_cache, gen_seq_len, gen_cu_q = specs[1] + assert (ctx_rows.start, ctx_rows.stop) == (0, 5) + assert (gen_rows.start, gen_rows.stop) == (5, 11) + assert ctx_cache.tolist() == [102, 203] + assert gen_cache.tolist() == [302, 401, 503] + assert (ctx_seq_len, gen_seq_len) == (0, 0) + assert ctx_cu_q is ctx_prefix + assert gen_cu_q.tolist() == [0, 2, 3, 6] + assert ctx_cu_q.data_ptr() != gen_cu_q.data_ptr() + + +def test_context_only_keeps_its_original_prefix_boundary(): + ctx_prefix = torch.tensor([0, 2, 5], dtype=torch.int32) + metadata = SimpleNamespace( + kv_lens_cuda_runtime=torch.tensor([102, 203], dtype=torch.int32), + num_ctx_tokens=5, + num_tokens=5, + max_seq_len=203, + mla_prepare_ctx_cu_seqlens=lambda: ctx_prefix, + ) + _, specs = _fused_q_rope_specs(SimpleNamespace(mqa=_Mqa()), metadata, 2, 0) + assert len(specs) == 1 + rows, cache_lens, seq_len, cu_q_seqlens = specs[0] + assert (rows.start, rows.stop) == (0, 5) + assert cache_lens.tolist() == [102, 203] + assert seq_len == 0 + assert cu_q_seqlens is ctx_prefix + + +def test_ragged_generation_fails_closed_without_stable_prefix_hook(): + metadata = SimpleNamespace( + kv_lens_cuda_runtime=torch.tensor([10, 20], dtype=torch.int32), + num_ctx_tokens=0, + num_tokens=3, + max_seq_len=20, + is_ragged_verify=True, + ) + mla = SimpleNamespace(mqa=_Mqa()) + assert _fused_q_rope_specs(mla, metadata, 0, 2) == (None, []) + metadata.mla_prepare_fused_q_gen_cu_seqlens = lambda: None + assert _fused_q_rope_specs(mla, metadata, 0, 2) == (None, []) + + +def test_prefix_storage_address_is_fixed_across_rebuilds(): + metadata = _ragged_metadata([2, 1, 3]) + first = metadata.mla_prepare_fused_q_gen_cu_seqlens() + first_address = first.data_ptr() + assert first.tolist() == [0, 2, 3, 6] + + metadata.seq_lens_cuda.copy_(torch.tensor([1, 4, 1], dtype=torch.int32)) + metadata._fused_q_gen_cu_seqlens_valid = False + second = metadata.mla_prepare_fused_q_gen_cu_seqlens() + assert second.data_ptr() == first_address + assert second.tolist() == [0, 1, 5, 6] + + +@pytest.mark.parametrize("capacity", [None, 2]) +def test_prefix_builder_fails_closed_without_graph_storage(capacity): + metadata = _ragged_metadata([1, 2]) + metadata.fused_q_gen_cu_seqlens = ( + None if capacity is None else torch.zeros(capacity, dtype=torch.int32) + ) + assert metadata.mla_prepare_fused_q_gen_cu_seqlens() is None + + +def test_sequence_update_invalidates_cached_prefix(): + metadata = DeepseekV4TrtllmAttentionMetadata.__new__(DeepseekV4TrtllmAttentionMetadata) + metadata._fused_q_gen_cu_seqlens_valid = True + metadata._invalidate_mla_scheduler_buffers() + assert not metadata._fused_q_gen_cu_seqlens_valid 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..a80f5d7f5b21 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py @@ -0,0 +1,305 @@ +# 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.backends.sparse.deepseek_v4.metadata import ( + DeepseekV4TrtllmAttentionMetadata, +) +from tensorrt_llm._torch.attention.backends.sparse.dsa import DSAtrtllmAttentionMetadata + + +class _StubCacheManager: + """Only the attributes the row expansion reads.""" + + def __init__(self, num_pools: int, max_blocks_per_seq: int) -> None: + 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: int, + max_blocks: int, + max_seqs: int, + draft: int = 5, + enable_ragged_verification: bool = True, +) -> DSAtrtllmAttentionMetadata: + """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.enable_ragged_verification = enable_ragged_verification + 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 test_feature_off_skips_ragged_row_buffers() -> None: + """Static DSA does not pay for confidence-only row metadata.""" + torch.cuda.init() + meta = _metadata( + num_pools=2, + max_blocks=8, + max_seqs=4, + enable_ragged_verification=False, + ) + + assert meta.gen_token_repeats_cuda is None + assert meta.row_kv_lens_cuda is None + assert meta.row_req_idx_cuda is None + assert meta.attn_row_kv_lens_cuda is None + assert meta.attn_row_block_offsets is None + + +@pytest.mark.parametrize( + "metadata_type", [DSAtrtllmAttentionMetadata, DeepseekV4TrtllmAttentionMetadata] +) +def test_feature_off_rejects_ragged_layout(metadata_type) -> None: + """An unexpected ragged update fails instead of dereferencing absent buffers.""" + meta = metadata_type.__new__(metadata_type) + meta.enable_ragged_verification = False + meta.ragged_verify_lens = None + assert not meta.is_ragged_verify + meta._validate_ragged_verification() + # A layout can be installed after construction, before the next prepare. + meta.ragged_verify_lens = [2] + assert meta.is_ragged_verify + + with pytest.raises(RuntimeError, match="without enabling ragged-verification buffers"): + meta.prepare() + with pytest.raises(RuntimeError, match="without enabling ragged-verification buffers"): + meta.apply_device_ragged_layout( + torch.tensor([2]), torch.tensor([0, 0]), torch.tensor([-1, 0]) + ) + + +@pytest.mark.parametrize( + "metadata_type", [DSAtrtllmAttentionMetadata, DeepseekV4TrtllmAttentionMetadata] +) +def test_device_window_kv_bound_preserves_context_and_input(metadata_type) -> None: + meta = metadata_type.__new__(metadata_type) + meta.device_windows_mode = True + meta.ragged_verify_lens = [2, 4] + meta._num_contexts = 1 + meta._seq_lens = torch.tensor([7, 2, 4], dtype=torch.int32) + meta._seq_lens_kv = None + meta.max_draft_tokens = 5 + meta.kv_lens = torch.tensor([7, 102, 204], dtype=torch.int32) + original = meta.kv_lens.clone() + bounded = meta._prepare_device_window_kv_lens(original) + assert original.tolist() == [7, 102, 204] + assert bounded.tolist() == [7, 106, 206] + assert meta.kv_lens.tolist() == bounded.tolist() + + +@pytest.mark.parametrize( + "device_windows_mode,ragged_lens", [(False, [2]), (True, None), (False, None)] +) +def test_native_kv_bound_reuses_storage(device_windows_mode, ragged_lens) -> None: + meta = DSAtrtllmAttentionMetadata.__new__(DSAtrtllmAttentionMetadata) + meta.device_windows_mode = device_windows_mode + meta.ragged_verify_lens = ragged_lens + lengths = torch.tensor([102], dtype=torch.int32) + assert meta._prepare_device_window_kv_lens(lengths) is lengths + + +def _expand( + meta: DSAtrtllmAttentionMetadata, + *, + num_contexts: int, + row_req_idx: list[int], + num_seqs: int, + block_offsets: torch.Tensor, +) -> int: + """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: int) -> None: + """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() -> None: + """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 {src.shape[1]} long" + ) + + +def test_context_rows_precede_generation_rows() -> None: + """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_device_layout_preserves_mixed_context_row_and_token_prefixes() -> None: + """Device-selected generation windows start after both context prefixes.""" + torch.cuda.init() + meta = _metadata(num_pools=1, max_blocks=8, max_seqs=4) + block_offsets = torch.arange(1 * 4 * 2 * 8, dtype=torch.int32, device="cuda").reshape( + 1, 4, 2, 8 + ) + rows = _expand( + meta, + num_contexts=2, + row_req_idx=[0, 1, 1], + num_seqs=4, + block_offsets=block_offsets, + ) + assert rows == 5 + + meta.num_generations = 2 + meta._num_ctx_tokens = 5 + meta.seq_lens = torch.tensor([2, 3, 1, 2], dtype=torch.int32) + meta._seq_lens_cuda = meta.seq_lens.to(device="cuda") + meta.req_idx_per_token = torch.full((32,), -1, dtype=torch.int32, device="cuda") + context_token_requests = torch.tensor([0, 0, 1, 1, 1], dtype=torch.int32, device="cuda") + meta.req_idx_per_token[:5].copy_(context_token_requests) + context_row_requests = torch.tensor([0, 1], dtype=torch.long, device="cuda") + context_row_corrections = torch.tensor([10, 20], dtype=torch.int32, device="cuda") + meta.attn_row_req_idx_cuda[:2].copy_(context_row_requests) + meta.attn_row_kv_correction_cuda[:2].copy_(context_row_corrections) + meta.indexer_k_cache_block_offsets = torch.arange( + 4 * 8, dtype=torch.int32, device="cuda" + ).reshape(4, 8) + + meta.apply_device_ragged_layout( + verify_lens=torch.tensor([1, 2], dtype=torch.int32, device="cuda"), + req_idx=torch.tensor([0, 1, 1], dtype=torch.long, device="cuda"), + kv_correction=torch.tensor([-2, -1, 0], dtype=torch.int32, device="cuda"), + ) + + assert torch.equal(meta.attn_row_req_idx_cuda[:5], torch.tensor([0, 1, 2, 3, 3], device="cuda")) + assert torch.equal( + meta.attn_row_kv_correction_cuda[:5], + torch.tensor([10, 20, -2, -1, 0], dtype=torch.int32, device="cuda"), + ) + assert torch.equal(meta.req_idx_per_token[:5], context_token_requests) + assert torch.equal( + meta.req_idx_per_token[5:8], + torch.tensor([2, 3, 3], dtype=torch.int32, device="cuda"), + ) + + +def test_every_runtime_view_shares_the_row_count() -> None: + """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/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index a5b81c528b5a..9f1ba6b50c6d 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -344,6 +344,7 @@ def make_mock(num_generations, kv_lens_list): num_generations=num_generations, num_sms=num_sms, max_draft_tokens=next_n - 1, + gen_token_stride=next_n, num_contexts=0, num_seqs=num_generations, kv_lens_cuda=kv_lens_cuda, @@ -473,6 +474,8 @@ def test_shared_topk_lifecycle(monkeypatch): ) metadata._create_kv_lens_2d_buffer = Mock() metadata.create_expanded_buffers = Mock() + metadata._ragged_num_rows = 0 + metadata._attn_num_rows = 0 with patch( "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata.prefer_pinned", @@ -548,6 +551,34 @@ def _predict(backend, input_type): assert metadata.shared_topk_indices is buffer +@pytest.mark.parametrize("enabled", [False, True]) +def test_kv_len_update_refreshes_ragged_rows_only_when_enabled(enabled): + metadata = DSAtrtllmAttentionMetadata.__new__(DSAtrtllmAttentionMetadata) + metadata.kv_cache_manager = None + metadata._num_tokens = 0 + metadata._num_generations = 0 + metadata.enable_ragged_verification = enabled + metadata._invalidate_pool_view_cache = Mock() + metadata._compute_kv_lens_row_reorder = Mock() + metadata.prepare_dense_topk_indices = Mock() + metadata.kv_lens_cuda = None + metadata.refresh_ragged_row_kv_lens = Mock() + metadata.refresh_token_major_gen_rows = Mock() + + with ( + patch.object(TrtllmAttentionMetadata, "on_update_kv_lens"), + patch( + "tensorrt_llm._torch.attention.backends.sparse.dsa.metadata._fused_dsa_meta_enabled", + return_value=False, + ), + ): + metadata.on_update_kv_lens() + + expected_calls = int(enabled) + assert metadata.refresh_ragged_row_kv_lens.call_count == expected_calls + assert metadata.refresh_token_major_gen_rows.call_count == expected_calls + + def test_indexer_post_load_weights_caches_fused_weight(): indexer = Indexer.__new__(Indexer) torch.nn.Module.__init__(indexer) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py index f5d9adcbc977..8ccaf29d749f 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_indexer_gvr_prior.py @@ -84,6 +84,8 @@ def test_indexer_forward_uses_prior_only_for_temporal_gvr( seq_lens=torch.tensor(([2, 3] if num_contexts else []) + [next_n] * num_generations), cuda_graph_buffers={}, is_cuda_graph=False, + is_ragged_verify=False, + ragged_row_kv_lens=Mock(return_value=None), get_empty=lambda buffers, shape, **kwargs: torch.full(shape, -1, dtype=torch.int32), skip_indexer_for_ctx_reqs=False, skip_indexer_for_gen_reqs=False, @@ -128,6 +130,7 @@ def decode( scores: torch.Tensor, sequence_lengths: torch.Tensor, scan_lengths: torch.Tensor, + row_kv_lens: torch.Tensor | None, output: torch.Tensor, n: int, max_seq_len: int, @@ -144,6 +147,7 @@ def decode( assert torch.all(prior == -99) else: assert prior is None + assert row_kv_lens is None assert n == next_n output.copy_(selections(num_gen_tokens)) return output diff --git a/tests/unittest/_torch/attention/test_fmha_interface.py b/tests/unittest/_torch/attention/test_fmha_interface.py index a265c3a69037..015b7497e734 100644 --- a/tests/unittest/_torch/attention/test_fmha_interface.py +++ b/tests/unittest/_torch/attention/test_fmha_interface.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import nullcontext from typing import cast from unittest.mock import Mock, patch @@ -85,6 +86,8 @@ def test_mla_forward_clears_fused_qkv_before_fmha_selection( attn._fmha_manager.select.return_value = fmha metadata = Mock(spec=TrtllmAttentionMetadata) + metadata.token_major_gen_view.return_value = None + metadata.presented_token_major.return_value = nullcontext() metadata.is_cross = False metadata.enable_flash_mla = False metadata.spec_bl_tree_first_sparse_mask_offset_kv = None @@ -112,6 +115,7 @@ def test_mla_forward_clears_fused_qkv_before_fmha_selection( output = TrtllmAttention.forward(attn, q, k, v, metadata, forward_args) + metadata.presented_token_major.assert_called_once_with(None) attn._fmha_manager.select.assert_called_once_with(attn, q, k, v, metadata, forward_args) fmha.forward.assert_called_once_with(q, k, v, metadata, forward_args) assert not forward_args.is_fused_qkv diff --git a/tests/unittest/_torch/modules/test_top_k.py b/tests/unittest/_torch/modules/test_top_k.py index c7f49d7ca692..6c306bb0e8f9 100644 --- a/tests/unittest/_torch/modules/test_top_k.py +++ b/tests/unittest/_torch/modules/test_top_k.py @@ -8,6 +8,7 @@ import pytest import torch +from torch._subclasses.fake_tensor import FakeTensorMode from tensorrt_llm._torch.modules.top_k import _CUTE_DSL_PREFILL_COPY_BITS, TopK, TopKImplementation @@ -118,6 +119,103 @@ def test_cute_dsl_radix_preserves_compressed_mtp_fallback(monkeypatch) -> None: assert runtime_call.kwargs["radix_aux_logits"].data_ptr() == radix_values.data_ptr() +@pytest.mark.parametrize("implementation", list(TopKImplementation)) +def test_ragged_decode_preserves_caller_radix_workspace(monkeypatch, implementation) -> None: + radix = Mock() + fallback_log = Mock() + monkeypatch.setattr("tensorrt_llm._torch.modules.top_k.logger.info_once", fallback_log) + monkeypatch.setattr(torch.ops.trtllm, "indexer_topk_decode", radix) + scores = torch.randn(3, 16) + output = torch.empty(3, 2, dtype=torch.int32) + logical = torch.tensor([16, 12], dtype=torch.int32) + scan = torch.tensor([4, 3], dtype=torch.int32) + row_kv_lens = torch.tensor([15, 16, 12], dtype=torch.int32) + indices = torch.empty(3, 10, 2, dtype=torch.int32) + values = torch.empty(3, 10, 2) + result = TopK(2, decode_implementation=implementation, compress_ratio=4)( + scores, + output, + is_prefill=False, + sequence_lengths=logical, + scan_lengths=scan, + row_kv_lens=row_kv_lens, + next_n=2, + radix_aux_indices=indices, + radix_aux_logits=values, + ) + assert result is output + radix.assert_called_once() + call = radix.call_args + assert call.kwargs["radix_aux_indices"].data_ptr() == indices.data_ptr() + assert call.kwargs["radix_aux_logits"].data_ptr() == values.data_ptr() + assert call.kwargs["row_kv_lens"] is row_kv_lens + assert (call.kwargs["row_kv_lens"] // 4).tolist() == [3, 4, 3] + if implementation == TopKImplementation.CUDA_RADIX: + fallback_log.assert_not_called() + else: + fallback_log.assert_called_once_with( + "Ragged per-row KV lengths require CUDA radix decode Top-K; " + f"using it instead of {implementation.value}.", + key="ragged_decode_radix_fallback", + ) + + +@pytest.mark.parametrize("ragged", [False, True]) +@pytest.mark.parametrize("compiled", [False, True]) +def test_indexer_topk_decode_fake_accepts_optional_row_kv_lens( + ragged: bool, compiled: bool +) -> None: + """Uniform and ragged decode must bind the fake in eager and compiled calls.""" + decode = torch.ops.trtllm.indexer_topk_decode + if compiled: + decode = torch.compile(decode, backend="eager", fullgraph=True) + with FakeTensorMode(): + scores = torch.empty(3, 16) + lengths = torch.empty(2, dtype=torch.int32) + output = torch.empty(3, 2, dtype=torch.int32) + row_kv_lens = torch.empty(3, dtype=torch.int32) if ragged else None + result = decode( + scores, + lengths, + output, + 1, + 2, + compress_ratio=4, + radix_aux_indices=torch.empty(3, 10, 2, dtype=torch.int32), + radix_aux_logits=torch.empty(3, 10, 2), + row_kv_lens=row_kv_lens, + ) + assert result is None + assert output.shape == (3, 2) and output.dtype == torch.int32 + + +@pytest.mark.parametrize("implementation", list(TopKImplementation)) +def test_compiled_ragged_decode_uses_radix_without_logging(monkeypatch, implementation) -> None: + """Fallback logging must not introduce a graph break during fake propagation.""" + fallback_log = Mock() + monkeypatch.setattr("tensorrt_llm._torch.modules.top_k.logger.info_once", fallback_log) + module = torch.compile( + TopK(2, decode_implementation=implementation, compress_ratio=4), + backend="eager", + fullgraph=True, + ) + with FakeTensorMode(): + output = torch.empty(3, 2, dtype=torch.int32) + result = module( + torch.empty(3, 16), + output, + is_prefill=False, + sequence_lengths=torch.empty(2, dtype=torch.int32), + scan_lengths=torch.empty(2, dtype=torch.int32), + row_kv_lens=torch.empty(3, dtype=torch.int32), + next_n=2, + radix_aux_indices=torch.empty(3, 10, 2, dtype=torch.int32), + radix_aux_logits=torch.empty(3, 10, 2), + ) + assert result is output + fallback_log.assert_not_called() + + def test_gvr_uses_caller_prior_state(monkeypatch) -> None: gvr = Mock() monkeypatch.setattr(torch.ops.trtllm, "cute_dsl_gvr_topk_decode", gvr) 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..e1e72fff16f2 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py @@ -0,0 +1,361 @@ +# 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 uses ``next_n`` to size its launch. Under DSpark's +ragged verification, the optional ``new_tokens_per_seq`` supplies the actual +per-request counts used for starting positions and compression boundaries. + +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 + +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 + + +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, + compress_ratio: int = 4, + ) -> None: + if len(kv_lens) != len(new_tokens): + raise ValueError("kv_lens and new_tokens must have the same length") + self.batch_size = len(kv_lens) + self.kv_lens = list(kv_lens) + self.new_tokens = list(new_tokens) + self.device = device + self.dtype = dtype + self.compress_ratio = compress_ratio + # CR4 overlaps adjacent windows; CR128 uses one head-width state. + state_dim = 2 * HEAD_DIM if compress_ratio == 4 else HEAD_DIM + + 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), 2 * state_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: torch.Tensor | None + ) -> tuple[torch.Tensor, 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, + self.compress_ratio, + next_n, + new_tokens_per_seq, + ) + return output, paged_kv, paged_score + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("compress_ratio", [4, 128]) +@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), + # CR128 distributes these windows over multiple chunk workers. + ([512, 768, 1024], 257), + ], +) +def test_uniform_new_tokens_matches_none( + kv_lens: list[int], next_n: int, dtype: torch.dtype, compress_ratio: int +) -> None: + """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, compress_ratio=compress_ratio + ) + + reference = case.run(next_n, None) + uniform_vector = torch.tensor(new_tokens, dtype=torch.int32, device=device) + ragged = case.run(next_n, uniform_vector) + + reference_output, reference_kv, reference_score = reference + ragged_output, ragged_kv, ragged_score = ragged + assert not torch.isnan(reference_output[: case.total_outputs]).any(), ( + "reference call left output rows unwritten; the test's cu_kv_comp does " + "not match what the kernel produces" + ) + assert torch.equal( + reference_output[: case.total_outputs], ragged_output[: case.total_outputs] + ), ( + "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_output.float() - ragged_output.float()).abs().max().item()}" + ) + assert torch.equal(reference_kv, ragged_kv) + assert torch.equal(reference_score, ragged_score) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize( + "compress_ratio,kv_lens,new_tokens", + [ + (4, [256, 300, 512, 1024], [6, 3, 1, 4]), + (128, [512, 512, 1024, 1024], [0, 3, 129, 257]), + ], +) +def test_ragged_new_tokens_matches_per_request_uniform( + dtype: torch.dtype, compress_ratio: int, kv_lens: list[int], new_tokens: list[int] +) -> None: + """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") + next_n = max(new_tokens) + + batched = _CompressorCase( + kv_lens, new_tokens, device, dtype, seed=99, compress_ratio=compress_ratio + ) + ragged_out, ragged_kv, ragged_score = 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, compress_ratio=compress_ratio + ) + # 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() + + start, end = cu_kv_comp[index], cu_kv_comp[index + 1] + if count == 0: + assert start == end + assert torch.equal(single.paged_kv, ragged_kv[block_start : block_start + num_blocks]) + assert torch.equal( + single.paged_score, ragged_score[block_start : block_start + num_blocks] + ) + continue + expected, expected_kv, expected_score = single.run(count, None) + expected = expected[: end - start] + 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()}" + ) + assert torch.equal(expected_kv, ragged_kv[block_start : block_start + num_blocks]) + assert torch.equal(expected_score, ragged_score[block_start : block_start + num_blocks]) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_ragged_append_without_output_updates_paged_state(dtype: torch.dtype) -> None: + """A request crossing no compression boundary still updates paged state.""" + device = torch.device("cuda") + kv_lens = [257, 260] + new_tokens = [1, 4] + batched = _CompressorCase(kv_lens, new_tokens, device, dtype, seed=7) + + ragged_out, ragged_kv, ragged_score = batched.run( + max(new_tokens), torch.tensor(new_tokens, dtype=torch.int32, device=device) + ) + assert batched.cu_kv_comp.tolist() == [0, 0, 1] + assert not torch.isnan(ragged_out[:1]).any() + + for index, (kv_len, count) in enumerate(zip(kv_lens, new_tokens)): + single = _CompressorCase([kv_len], [count], device, dtype, seed=7) + 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] + block_start = int(block_row.min().item()) + num_blocks = block_row.shape[1] + single.block_table = block_row - block_start + 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_out, expected_kv, expected_score = single.run(count, None) + start, end = batched.cu_kv_comp[index : index + 2].tolist() + assert torch.equal(expected_out[: end - start], ragged_out[start:end]) + assert torch.equal(expected_kv, ragged_kv[block_start : block_start + num_blocks]) + assert torch.equal(expected_score, ragged_score[block_start : block_start + num_blocks]) + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +def test_ragged_cuda_graph_replay_matches_eager(compress_ratio: int) -> None: + """Captured launches keep the per-request counts and paged writes intact.""" + device = torch.device("cuda") + counts = [0, 3, 129, 257] + case = _CompressorCase( + [512, 512, 1024, 1024], + counts, + device, + torch.bfloat16, + seed=81, + compress_ratio=compress_ratio, + ) + vector = torch.tensor(counts, dtype=torch.int32, device=device) + expected = case.run(max(counts), vector) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + case.run(max(counts), vector) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = case.run(max(counts), vector) + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + for reference, replayed in zip(expected, actual): + assert torch.equal(reference, replayed) 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..ac78a086f721 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py @@ -0,0 +1,344 @@ +# 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 +MAX_RADIX_BLOCKS_PER_ROW = 10 # kMaxBlocksPerRowDecode in indexerTopK.cu. + +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, + index_topk: int = INDEX_TOPK, +) -> 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 * MAX_RADIX_BLOCKS_PER_ROW * index_topk + 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) + + +@pytest.mark.parametrize( + ("dtype", "index_topk", "num_columns"), + [(torch.bfloat16, 512, 8192), (torch.float32, 2048, 262144)], +) +def test_ragged_dtype_and_forced_split_match_torch(dtype, index_topk, num_columns): + """Cover BF16 ragged dispatch and the FP32 forced ten-block split tier.""" + verify_lens = [2, 1] + kv_lens = [num_columns - 3, num_columns] + row_kv_lens = _row_kv_lens(kv_lens, verify_lens) + logits = _make_logits(sum(verify_lens), num_columns, seed=42).to(dtype) + actual = _run(logits, kv_lens, 1, row_kv_lens=row_kv_lens, index_topk=index_topk) + + valid = torch.arange(num_columns, device="cuda").unsqueeze(0) < row_kv_lens.unsqueeze(1) + reference = logits.masked_fill(~valid, float("-inf")).topk(index_topk, dim=-1).values + assert torch.all((actual >= 0) & (actual < row_kv_lens.unsqueeze(1))) + assert torch.all(actual.sort(dim=-1).values.diff(dim=-1) > 0) + selected = logits.gather(1, actual.long()).sort(dim=-1).values + # BF16 can tie at the cutoff: exact selected values, not tie-breaking order, + # define equivalence. Extents and uniqueness are checked independently. + torch.testing.assert_close(selected, reference.sort(dim=-1).values, atol=0, rtol=0) + + +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"), + ) + + +@pytest.mark.parametrize( + ("compress_ratio", "invalid_extent"), + [ + (1, -1), + (1, NUM_COLUMNS + 1), + (4, (NUM_COLUMNS + 1) * 4), + ], +) +def test_invalid_row_kv_extent_fails_closed(compress_ratio: int, invalid_extent: int) -> None: + """Device-selected extents must never address outside a logits row. + + Two rows of width 8192 select four blocks per row on supported GPUs, so + this also pins the fail-closed behavior through the split/merge path. + """ + logits = _make_logits(2, NUM_COLUMNS, seed=5) + row_kv_lens = torch.tensor([KV_LENS[0], invalid_extent], dtype=torch.int32, device="cuda") + + actual = _run( + logits, + KV_LENS[:2], + next_n=1, + row_kv_lens=row_kv_lens, + compress_ratio=compress_ratio, + ) + + assert torch.all(actual[1] == -1) + + +@pytest.mark.parametrize("compress_ratio", [1, 4]) +def test_full_width_row_kv_extent_is_valid(compress_ratio: int) -> None: + """The exact logical row width is valid, not an out-of-range extent.""" + logits = _make_logits(2, NUM_COLUMNS, seed=6) + full_width_extent = NUM_COLUMNS * compress_ratio + seq_lens = [full_width_extent, full_width_extent] + + uniform = _run(logits, seq_lens, next_n=1, compress_ratio=compress_ratio) + ragged = _run( + logits, + seq_lens, + next_n=1, + row_kv_lens=torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + compress_ratio=compress_ratio, + ) + + _assert_same_topk(ragged, uniform) + + +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 * MAX_RADIX_BLOCKS_PER_ROW * 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)