[None][feat] Kimi K3 NVFP4 on Rubin (SM107): Modeling V2 target and the runtime it needs

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,18 +180,24 @@ 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 both servers enable compatible speculation and shard
+attention alike: the same `tensor_parallel_size`, where `enable_attention_dp`
+counts as 1 because every rank keeps all KV heads.
 
 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) keep
+a standalone DSpark MLA drafter's layers after their own, with the CUTEDSL,
+TRTLLM or VANILLA drafter backend.
 
 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

tensorrt_llm/_torch/_experimental/modeling_v2/models/kimi_k3_vl/kimi_k3_nvfp4__sm_107__default/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Kimi K3 (NVFP4) / sm_107 / default.
+
+Imports nothing: ``routing.py`` selects the deployment's ``features`` profile before any module
+of the target is imported.
+"""