[None][feat] Staircase: GLM-5.2 NVFP4 low-latency decode target for Rubin (modeling_v2, TP16)

3rdparty/fetch_content.json

@@ -46,7 +46,10 @@
       "git_tag": "b64107f2b9599ca76445b7f62eedf66bae1d095b",
       "git_submodules_recurse": true,
       "source_subdir": "dont-add-this-project-with-add-subdirectory",
-      "patch_file": "patches/deepgemm_configurable_barrier_timeout.patch"
+      "patch_file": [
+        "patches/deepgemm_configurable_barrier_timeout.patch",
+        "patches/deepgemm_sm10x_family_target.patch"
+      ]
     },
     {
       "name": "eigen",

3rdparty/patches/deepgemm_sm10x_family_target.patch

@@ -0,0 +1,13 @@
+--- a/csrc/runtime/jit.hpp
++++ b/csrc/runtime/jit.hpp
+@@ -37,6 +37,10 @@
+         // Preserve the existing 12.x family target when DeepJIT derives a different minor target.
+         if (runtime->device.get_arch_major() == 12 and runtime->device.get_arch_minor() != 0)
+             runtime->default_compiler_options.arch = "120f";
++        // A 10.x device past 10.3 (e.g. 10.7) builds for the 10.0 family target, which it runs:
++        // the bundled CUTLASS enables TMA and tcgen05 only for the targets it knows.
++        if (runtime->device.get_arch_major() == 10 and runtime->device.get_arch_minor() > 3)
++            runtime->default_compiler_options.arch = "100f";
+         return runtime;
+     });
+ }

cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu

@@ -184,19 +184,20 @@ CUresult copyDeviceToDevice(std::vector<MMTask> const& tasks, ssize_t numBytes,
     return launchBatchedCopy(false, tasks, numBytes, stream);
 }
 
-// dst_tensor[:, :num_seqs, 0] = src_tensor[:, copy_idx]
-// dst_tensor[:, :num_seqs, 1] = dst_tensor[:, :num_seqs, 0] + 1
+// dst_tensor[:, :num_seqs, 0, :numBlocksToCopy] = src_tensor[:, copy_idx, 0, :numBlocksToCopy]
+// dst_tensor[:, :num_seqs, 1, :numBlocksToCopy] = dst_tensor[:, :num_seqs, 0, :numBlocksToCopy] + 1
+// Rows are numBlocksPerSeq long in both tensors; the columns past numBlocksToCopy are not written.
 __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,
-    SizeType32 const* __restrict__ kvOffset)
+    SizeType32 numBlocksPerSeq, SizeType32 numBlocksToCopy, SizeType32 const* __restrict__ copyIndex,
+    SizeType32 const* __restrict__ indexScales, SizeType32 const* __restrict__ kvOffset)
 {
     constexpr uint32_t kvFactor = 2;
     constexpr auto elemPerAccess = sizeof(PackedInt) / sizeof(SizeType32);
 
     __shared__ PackedInt data[copyBlocknbBufs][copyBlockCtaSize];
 
-    auto const iterPerSeq = divUp(numBlocksPerSeq * sizeof(SizeType32), sizeof(PackedInt) * copyBlockCtaSize);
+    auto const iterPerSeq = divUp(numBlocksToCopy * sizeof(SizeType32), sizeof(PackedInt) * copyBlockCtaSize);
     auto const tid = threadIdx.x;
     auto const poolIdx = blockIdx.x;
     auto const seqIdx = blockIdx.y;
@@ -205,7 +206,7 @@ __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict
     uint32_t const dstIdxKBeg = tid * elemPerAccess + (poolIdx * dstMaxNumSequences + seqIdx) * seqDimStride;
     uint32_t const dstIdxVBeg = dstIdxKBeg + numBlocksPerSeq;
 
-    uint32_t const srcIdxEnd = (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride + numBlocksPerSeq;
+    uint32_t const srcIdxEnd = (poolIdx * srcMaxNumSequences + copyIndex[seqIdx]) * seqDimStride + numBlocksToCopy;
 
     for (uint32_t i = 0; i < iterPerSeq + copyBlocknbBufs; i++)
     {
@@ -250,7 +251,7 @@ __global__ void copyBatchBlockOffsetsToDeviceKernel(SizeType32 const* __restrict
 
 // Host-side launcher
 void copyBatchBlockOffsetsToDevice(ITensor const& input, ITensor& output, ITensor const& copyIndex,
-    ITensor const& indexScales, ITensor const& kvOffset, CUstream stream) noexcept
+    ITensor const& indexScales, ITensor const& kvOffset, SizeType32 maxNumBlocks, CUstream stream) noexcept
 {
     using namespace tensorrt_llm::runtime;
 
@@ -282,11 +283,22 @@ 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)));
 
+    // A negative maxNumBlocks copies whole rows. Otherwise the copy covers the first maxNumBlocks
+    // columns, rounded up to whole PackedInt accesses, which stay inside the row.
+    SizeType32 numBlocksToCopy = numBlocksPerSeq;
+    if (maxNumBlocks >= 0)
+    {
+        constexpr SizeType32 blocksPerAccess = sizeof(PackedInt) / sizeof(SizeType32);
+        SizeType32 const roundedBlocks
+            = (std::max(maxNumBlocks, SizeType32{1}) + blocksPerAccess - 1) / blocksPerAccess * blocksPerAccess;
+        numBlocksToCopy = std::min(numBlocksPerSeq, roundedBlocks);
+    }
+
     dim3 gridDim(numPools, numSeqs, 1);
     dim3 blockDim(copyBlockCtaSize);
 
     copyBatchBlockOffsetsToDeviceKernel<<<gridDim, blockDim, 0, stream>>>(srcPtr, dstPtr, srcMaxNumSequences,
-        dstMaxNumSequences, numBlocksPerSeq, copyIndexPtr, indexScalesPtr, kvOffsetPtr);
+        dstMaxNumSequences, numBlocksPerSeq, numBlocksToCopy, copyIndexPtr, indexScalesPtr, kvOffsetPtr);
 }
 
 } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2

cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h

@@ -108,7 +108,9 @@ CUresult copyHostToDevice(std::vector<Task<MemAddress, MemAddress>> const& tasks
 CUresult copyDeviceToHost(std::vector<Task<MemAddress, MemAddress>> const& tasks, ssize_t numBytes, CUstream stream);
 CUresult copyDeviceToDevice(std::vector<Task<MemAddress, MemAddress>> const& tasks, ssize_t numBytes, CUstream stream);
 
+//! Copies the block offsets of the batch rows. A non-negative maxNumBlocks bounds the copied columns
+//! (rounded up to a multiple of four); a negative one copies whole rows.
 void copyBatchBlockOffsetsToDevice(ITensor const& input, ITensor& output, ITensor const& copyIndex,
-    ITensor const& indexScales, ITensor const& kvOffset, CUstream stream) noexcept;
+    ITensor const& indexScales, ITensor const& kvOffset, SizeType32 maxNumBlocks, CUstream stream) noexcept;
 
 } // namespace tensorrt_llm::batch_manager::kv_cache_manager_v2