[None][feat] Shared drafter KV pools, FP8 MLA context inputs and KV cache V2 copy/memory fixes for hybrid speculative models

cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu

@@ -30,6 +30,9 @@ namespace tensorrt_llm::batch_manager::kv_cache_manager_v2
 {
 constexpr uint32_t copyBlockCtaSize = 128;
 constexpr uint32_t copyBlocknbBufs = 2;
+// Loads per thread of one block-offset copy CTA: each CTA copies 1024 entries of a row,
+// so a long row is read from host memory by many CTAs in parallel.
+constexpr uint32_t copyBlockItersPerCta = 2;
 
 __device__ __host__ inline uint32_t divUp(uint32_t a, uint32_t b)
 {
@@ -78,6 +81,7 @@ void copyBasePageRowsToDevice(ITensor const& input, ITensor& output, SizeType32
 
 // dst_tensor[:, :num_seqs, 0] = src_tensor[:, copy_idx]
 // dst_tensor[:, :num_seqs, 1] = dst_tensor[:, :num_seqs, 0] + 1
+// Grid: (pool, sequence, chunk of copyBlockItersPerCta loads per thread along the row).
 __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict__ srcPtr,
     SizeType32* __restrict__ dstPtr, SizeType32 const srcMaxNumSequences, SizeType32 const dstMaxNumSequences,
     SizeType32 numBlocksPerSeq, SizeType32 const* __restrict__ copyIndex, SizeType32 const* __restrict__ indexScales,
@@ -89,17 +93,21 @@ __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict
     __shared__ PackedInt data[copyBlocknbBufs][copyBlockCtaSize];
 
     auto const iterPerSeq = divUp(numBlocksPerSeq * sizeof(SizeType32), sizeof(PackedInt) * copyBlockCtaSize);
+    auto const iterBeg = blockIdx.z * copyBlockItersPerCta;
+    auto const numIters = min(copyBlockItersPerCta, iterPerSeq - iterBeg);
     auto const tid = threadIdx.x;
     auto const poolIdx = blockIdx.x;
     auto const seqIdx = blockIdx.y;
     auto const seqDimStride = kvFactor * numBlocksPerSeq;
-    uint32_t const srcIdxBeg = tid * elemPerAccess + (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride;
-    uint32_t const dstIdxKBeg = tid * elemPerAccess + (poolIdx * dstMaxNumSequences + seqIdx) * seqDimStride;
+    uint32_t const chunkBeg = iterBeg * copyBlockCtaSize * elemPerAccess;
+    uint32_t const srcIdxBeg
+        = chunkBeg + tid * elemPerAccess + (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride;
+    uint32_t const dstIdxKBeg = chunkBeg + tid * elemPerAccess + (poolIdx * dstMaxNumSequences + seqIdx) * seqDimStride;
     uint32_t const dstIdxVBeg = dstIdxKBeg + numBlocksPerSeq;
 
     uint32_t const srcIdxEnd = (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride + numBlocksPerSeq;
 
-    for (uint32_t i = 0; i < iterPerSeq + copyBlocknbBufs; i++)
+    for (uint32_t i = 0; i < numIters + copyBlocknbBufs; i++)
     {
         uint32_t const idxBuf = i % copyBlocknbBufs;
         if (i >= copyBlocknbBufs)
@@ -129,7 +137,7 @@ __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict
         PackedInt* const dst = &data[idxBuf][tid];
         uint32_t const srcIdx = srcIdxBeg + copyBlockCtaSize * ldIter * elemPerAccess;
         PackedInt const* const src = reinterpret_cast<PackedInt const*>(srcPtr + srcIdx);
-        if (srcIdx < srcIdxEnd)
+        if (ldIter < numIters && srcIdx < srcIdxEnd)
         {
             uint32_t const size = sizeof(PackedInt);
             asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"l"(__cvta_generic_to_shared(dst)),
@@ -174,7 +182,16 @@ void copyBatchBlockOffsetsToDevice(ITensor const& input, ITensor& output, ITenso
         "Not implemented case: numBlocksPerSeq * sizeof(SizeType32) = %zu must be a multiple of %zu.",
         static_cast<size_t>(numBlocksPerSeq * sizeof(SizeType32)), static_cast<size_t>(sizeof(PackedInt)));
 
-    dim3 gridDim(numPools, numSeqs, 1);
+    auto const iterPerSeq = divUp(static_cast<uint32_t>(numBlocksPerSeq * sizeof(SizeType32)),
+        static_cast<uint32_t>(sizeof(PackedInt) * copyBlockCtaSize));
+    auto const numChunks = divUp(iterPerSeq, copyBlockItersPerCta);
+    if (numChunks == 0)
+    {
+        return;
+    }
+    TLLM_CHECK_WITH_INFO(
+        numChunks <= 65535, "numBlocksPerSeq = %d needs %u CTAs per row, more than 65535.", numBlocksPerSeq, numChunks);
+    dim3 gridDim(numPools, numSeqs, numChunks);
     dim3 blockDim(copyBlockCtaSize);
 
     copyBatchBlockOffsetsToDeviceKernel<<<gridDim, blockDim, 0, stream>>>(srcPtr, dstPtr, srcMaxNumSequences,

cpp/tensorrt_llm/kernels/mlaKernels.cu

@@ -1906,6 +1906,27 @@ void invokeMLAContextFp8Quantize(MlaParams<T>& params, int total_kv_len, cudaStr
                         params.quant_scale_o, params.dequant_scale_q, params.dequant_scale_kv, params.host_bmm1_scale);
             }
         }
+        else if (params.fp8_kv_input)
+        {
+            // K and V are FP8 already: quantize Q only, or nothing when it is FP8 too. Block 0 still
+            // writes the BMM scales.
+            constexpr int threads_per_block = 384;
+            constexpr int num_tokens_per_block = threads_per_block * 16 / 192 * sizeof(T);
+            int const q_len = params.fp8_q_input ? 0 : params.acc_q_len;
+            dim3 grid(std::max(1, int(tensorrt_llm::common::divUp(q_len, num_tokens_per_block))), 1,
+                params.fp8_q_input ? 1 : params.head_num);
+
+            TLLM_LOG_DEBUG(
+                "Launching quantizeCopyInputToFp8Kernel (FP8 K/V input) with grid_size: (%d, %d, %d), "
+                "threads_per_block: %d, acc_q_len: %d, fp8_q_input: %d",
+                grid.x, grid.y, grid.z, threads_per_block, params.acc_q_len, params.fp8_q_input);
+
+            quantizeCopyInputToFp8Kernel<T, threads_per_block, 128, 64, 128, false>
+                <<<grid, threads_per_block, 0, stream>>>(params.q_buf, static_cast<__nv_fp8_e4m3*>(params.quant_q_buf),
+                    nullptr, nullptr, nullptr, nullptr, q_len, 0, srcVTokenStride, params.quant_scale_qkv,
+                    params.bmm1_scale, params.bmm2_scale, params.quant_scale_o, params.dequant_scale_q,
+                    params.dequant_scale_kv, params.host_bmm1_scale);
+        }
         else
         {
             // The Q or K tensor has layout of [num_tokens, head_num, 192] in the non-absorption mode.

cpp/tensorrt_llm/kernels/mlaKernels.h

@@ -150,6 +150,14 @@ struct MlaParams
     // `quant_q_buf`. Context preprocessing must still rotate/cache K.
     bool q_rope_applied = false;
 
+    // Non-absorption context MLA: `k_buf` / `v_buf` point to contiguous FP8 E4M3 K
+    // [total_kv_len, head_num * (qk_nope_head_dim + qk_rope_head_dim)] and V
+    // [total_kv_len, head_num * v_head_dim], which the FMHA reads as they are.
+    bool fp8_kv_input = false;
+    // With `fp8_kv_input`: `quant_q_buf` already holds the FP8 Q, so the context
+    // quantization only writes the BMM scales.
+    bool fp8_q_input = false;
+
     // DSv4 fused inverse-RoPE + FP8 quant epilogue parameters.
     Dsv4EpilogueFusionParams dsv4_epilogue_fusion;
 

cpp/tensorrt_llm/thop/attentionOp.cpp

@@ -2881,6 +2881,34 @@ void MLA::prepareInputs(FmhaParams& p, bool isGen)
                 "Separate-K context FMHA requires contiguous K features within each token");
             TORCH_CHECK(p.value_input->strides()[1] == 1,
                 "Separate-V context FMHA requires contiguous V features within each token");
+            auto const& k = p.key_input.value();
+            auto const& v = p.value_input.value();
+            TORCH_CHECK(k.scalar_type() == v.scalar_type(), "MLA context k and v must have the same dtype.");
+            // FP8 k / v (and an FP8 Q in quant_q_buffer, if given with quant_scale_qkv): the caller already
+            // quantized them for the FP8 context FMHA, which reads them instead of quantizing BF16 inputs.
+            if (k.scalar_type() == torch::kFloat8_e4m3fn)
+            {
+                TORCH_CHECK(mFP8ContextMLA, "FP8 MLA context k / v require the FP8 context MLA FMHA.");
+                TORCH_CHECK(!p.fwd.latent_cache.has_value(),
+                    "FP8 MLA context k / v are for cached-KV context attention, without latent_cache.");
+                TORCH_CHECK(k.is_contiguous() && v.is_contiguous(), "FP8 MLA context k and v must be contiguous.");
+                TORCH_CHECK(k.size(1) == num_heads * (mMLAParams.qk_nope_head_dim + mMLAParams.qk_rope_head_dim)
+                        && v.size(1) == num_heads * mMLAParams.v_head_dim,
+                    "FP8 MLA context k / v must be [kv_tokens, num_heads * qk_head_dim] / [kv_tokens, num_heads * "
+                    "v_head_dim].");
+                if (p.fwd.quant_q_buffer.has_value() && p.fwd.quant_scale_qkv.has_value())
+                {
+                    auto const& quantQ = p.fwd.quant_q_buffer.value();
+                    TORCH_CHECK(quantQ.scalar_type() == torch::kFloat8_e4m3fn && quantQ.is_contiguous()
+                            && quantQ.numel() >= p.getQueryInput().numel(),
+                        "The FP8 MLA context quant_q_buffer must be a contiguous FP8 copy of q.");
+                }
+            }
+            else
+            {
+                TORCH_CHECK(k.scalar_type() == p.getQueryInput().scalar_type(),
+                    "MLA context k / v must be FP8 E4M3 or have the dtype of q.");
+            }
         }
     }
     else
@@ -3291,6 +3319,8 @@ int MLA::enqueueContext(FmhaParams const& p, cudaStream_t stream)
     {
         mla.k_buf = p.getK<T>();
         mla.v_buf = p.getV<T>();
+        mla.fp8_kv_input = p.key_input.has_value() && p.key_input->scalar_type() == torch::kFloat8_e4m3fn;
+        mla.fp8_q_input = mla.fp8_kv_input && p.fwd.quant_q_buffer.has_value() && p.fwd.quant_scale_qkv.has_value();
         mla.helix_position_offsets = p.getHelixPositionOffsets();
         mla.helix_is_inactive_rank = p.getHelixIsInactiveRank();
     }
@@ -3453,8 +3483,21 @@ int MLA::enqueueContext(FmhaParams const& p, cudaStream_t stream)
             "FP8 v buffer is required for fp8 context MLA in non-sparse mode");
 
         fmhaParams.qPtr = reinterpret_cast<void const*>(fp8QBuf);
-        fmhaParams.kPtr = useSparseMLA() ? nullptr : reinterpret_cast<void const*>(workspaceViews.fp8KBuf);
-        fmhaParams.vPtr = useSparseMLA() ? nullptr : reinterpret_cast<void const*>(workspaceViews.fp8VBuf);
+        if (useSparseMLA())
+        {
+            fmhaParams.kPtr = nullptr;
+            fmhaParams.vPtr = nullptr;
+        }
+        else if (mla.fp8_kv_input)
+        {
+            fmhaParams.kPtr = reinterpret_cast<void const*>(mla.k_buf);
+            fmhaParams.vPtr = reinterpret_cast<void const*>(mla.v_buf);
+        }
+        else
+        {
+            fmhaParams.kPtr = reinterpret_cast<void const*>(workspaceViews.fp8KBuf);
+            fmhaParams.vPtr = reinterpret_cast<void const*>(workspaceViews.fp8VBuf);
+        }
     }
     else
     {

docs/source/features/speculative-decoding.md

@@ -180,23 +180,37 @@ target pages. The draft layers keep the drafter's own dtype, so an FP8 target KV
 cache (`kv_cache_config.dtype: fp8`) is paired with a BF16 draft cache.
 
 In disaggregated serving, direct worker launches require
-`TRTLLM_DISAGG_SHARED_DRAFT_KV=1` on both the context and generation workers
-to enable shared draft-KV transfer. The `disaggregated_mpi_worker` launcher
-enables this setting automatically when both servers enable compatible
-speculation and use the same `tensor_parallel_size` and
-`enable_attention_dp`.
+`TRTLLM_DISAGG_SHARED_DRAFT_KV=1` (or `auto`, which shares wherever the layout
+allows it) on both the context and generation workers to enable shared
+draft-KV transfer. The `disaggregated_mpi_worker` launcher enables this setting
+automatically when its config lists both context and generation servers, both
+enable compatible speculation, and they shard attention alike: the same
+`tensor_parallel_size`, where `enable_attention_dp` counts as 1 because every
+rank keeps all KV heads. A config that lists only one side (for example, a
+context-only MPI world) leaves sharing off unless the variable is `1`.
 
 When sharing is enabled, the context server's draft pages are transferred to
 the generation server with the target pages, so the drafter starts generation
 with the prompt context instead of an empty one. This requires the draft layers
 to have the target's rank-local KV-head count; the draft dtype and head
 dimension may differ, except that an NVFP4 target KV cache requires an
-identical draft layout. Both servers must use the same layout.
+identical draft layout unless both caches are K-only (MLA). Both servers must
+use the same layout.
+
+Hybrid recurrent/attention targets whose cache manager opts in (Kimi K3) also
+hold a standalone DSpark drafter's layers, after their own. In every case the
+drafter decides whether it shares: it declares whether its block attention reads
+the manager's pages at the run's `tokens_per_block`
+(`DFlashForCausalLM.reads_managed_ctx_pages`). By default only
+`attention_backend="TRTLLM"` does, at the page sizes its trtllm-gen kernels
+support; a drafter whose other backends read paged context declares that itself.
 
 The drafter keeps a separate draft KV cache manager instead when any of the
-following applies: context parallelism larger than 1, a per-layer KV-head count
-on the target or drafter, a specialized KV cache manager (for example, hybrid Mamba
-models), Gemma 4 hybrid attention, or, in disaggregated serving, a draft layout
+following applies: context parallelism larger than 1 (helix context parallelism
+is shared only with a drafter that reads a helix-sharded context), a per-layer
+KV-head count on the target or drafter, a specialized KV cache manager that does
+not opt in (for example, most hybrid Mamba models), Gemma 4 hybrid attention, or,
+in disaggregated serving, a draft layout
 that cannot be transferred with the target's. When the layers are shared, the
 worker log reports `target and draft KV layers share one KV cache manager` at
 startup.