User/itabrizian/fix disagg gen init history length

.github/workflows/blossom-ci.yml

@@ -74,6 +74,7 @@ jobs:
         "bobboli",
         "Boreas618",
         "brb-nv",
+        "brnguyen2",
         "byshiue",
         "CarstyYou",
         "cascade812",
@@ -215,6 +216,7 @@ jobs:
         "milesial",
         "MinaHuai",
         "ming-wei",
+        "Mingyang",
         "mk-nvidia",
         "mlefeb01",
         "moraxu",
@@ -363,6 +365,7 @@ jobs:
         "yingguo-trt",
         "yiqingy0",
         "yizhang-nv",
+        "yongzhiz",
         "yuanjings-nvda",
         "yuanjingx87",
         "yuantailing",

.pre-commit-config.yaml

@@ -2726,8 +2726,8 @@ static-analysis-files: &static_analysis_files |
         tests/unittest/_torch/sampler/test_beam_search_util.py |
     )$
 
-# Global exclude: vendored code + trtllm-gen FMHA artifacts (cubin pointers, export headers, cuda_ptx)
-exclude: '(^triton_kernels/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/)'
+# Global exclude: vendored code + trtllm-gen generated artifacts (cubin pointers, export headers, cuda_ptx)
+exclude: '(^triton_kernels/|trtllmGenKernels/fmha/cubin/kernelMetaInfo\.h$|trtllmGenKernels/batchedGemm/trtllmGen_bmm_export/KernelMetaInfo\.h$|cubin\.cpp$|cubin\.h$|trtllmGenKernels/fmha/trtllmGen_fmha_export/|trtllmGenKernels/fmha/cuda_ptx/)'
 
 default_install_hook_types: [pre-commit, commit-msg]
 repos:

3rdparty/fetch_content.json

@@ -31,7 +31,7 @@
     {
       "name": "deepgemm",
       "git_repository": "https://github.com/deepseek-ai/DeepGEMM",
-      "git_tag": "c491439ed5966833d56883ca302b6f72e74f8105",
+      "git_tag": "67fc64863d43521080bf2005e6528d0fceee9510",
       "git_submodules_recurse": true,
       "source_subdir": "dont-add-this-project-with-add-subdirectory"
     },

ATTRIBUTIONS-Python.md

@@ -5261,7 +5261,7 @@ For more information, please refer to <http://unlicense.org>
   - `Tracker`: https://github.com/tox-dev/py-filelock/issues
 
 
-## flashinfer-python (0.6.10)
+## flashinfer-python (0.6.11)
 
 ### Licenses
 License: `Apache-2.0`
@@ -33239,7 +33239,7 @@ License: `NVIDIA Proprietary Software`
   - `Homepage`: https://developer.nvidia.com/cusparselt
 
 
-## nvidia-cutlass-dsl (4.4.2)
+## nvidia-cutlass-dsl (4.5.0)
 
 ### Licenses
 License: `None`

cpp/include/tensorrt_llm/batch_manager/llmRequest.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-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.
@@ -1835,6 +1835,15 @@ class GenericLlmRequest
         return mPerfMetrics.kvCacheMetrics.numNewAllocatedBlocks;
     }
 
+    void updateKvCachePerfMetrics(
+        SizeType32 allocTotalBlocks, SizeType32 allocNewBlocks, SizeType32 reusedBlocks, SizeType32 missedBlocks)
+    {
+        updateAllocTotalBlocksPerRequest(allocTotalBlocks);
+        updateAllocNewBlocksPerRequest(allocNewBlocks);
+        updateReusedBlocksPerRequest(reusedBlocks);
+        updateMissedBlocksPerRequest(missedBlocks);
+    }
+
     void updateReusedBlocksPerRequest(SizeType32 reusedBlocksPerRequest)
     {
         mPerfMetrics.kvCacheMetrics.numReusedBlocks += reusedBlocksPerRequest;

cpp/include/tensorrt_llm/executor/transferAgent.h

@@ -18,6 +18,7 @@
 
 #include "tensorrt_llm/common/assert.h"
 #include <fcntl.h>
+#include <map>
 #include <memory>
 #include <mutex>
 #include <optional>
@@ -194,8 +195,57 @@ using RegisterDescs = MemoryDescs;
 using SyncMessage = std::string;
 using ConnectionInfoType = std::string;
 
+/// Per-region VMM chunk info used for splitting descriptors at chunk boundaries.
+struct VramRegionInfo
+{
+    size_t totalLen;
+    size_t chunkSize; ///< 0 = cudaMalloc (no split), >0 = VMM chunk size
+};
+
+/// Region map: virtual base address → region info.
+using VramRegionMap = std::map<uintptr_t, VramRegionInfo>;
+
+/// Backend-agnostic VMM descriptor split utilities (no NIXL dependency).
+struct VmmDescSplitter
+{
+    /// @brief Look up VMM chunk info for an address from a region map.
+    /// @return {chunkSize, regionBase}. Returns {0, 0} if address is not in any region.
+    [[nodiscard]] static std::pair<size_t, uintptr_t> lookupChunkInfo(uintptr_t addr, VramRegionMap const& regionMap);
+
+    /// @brief Split VRAM descs at chunk boundaries using a pre-built region map.
+    /// For non-VRAM or addresses not in the map, descs pass through unchanged.
+    [[nodiscard]] static MemoryDescs splitDescsWithRegionMap(MemoryDescs const& descs, VramRegionMap const& regionMap);
+
+    /// @brief Split paired src/dst descs using local and remote region maps.
+    /// src is split by localRegionMap, dst is split by remoteRegionMap.
+    /// The final piece size is min(srcPiece, dstPiece, remaining).
+    [[nodiscard]] static std::pair<MemoryDescs, MemoryDescs> splitTransferDescsWithRegionMaps(
+        MemoryDescs const& srcDescs, MemoryDescs const& dstDescs, VramRegionMap const& localRegionMap,
+        VramRegionMap const& remoteRegionMap);
+
+    /// @brief Split VRAM descs at VMM chunk boundaries detected via cuMemGetAddressRange.
+    /// For cudaMalloc memory (single allocation), descs pass through unchanged.
+    /// @param[out] detectedChunkSize Set to the VMM chunk size if detected, 0 otherwise.
+    [[nodiscard]] static MemoryDescs splitVmmDescs(MemoryDescs const& descs, size_t& detectedChunkSize);
+
+    /// @brief Build a VramRegionMap by probing each VRAM descriptor with cuMemGetAddressRange.
+    /// For each descriptor, detects whether it spans multiple VMM chunks and records {totalLen, chunkSize}.
+    /// @param descs VRAM memory descriptors to probe.
+    /// @return Region map with per-descriptor VMM info (chunkSize=0 for cudaMalloc memory).
+    [[nodiscard]] static VramRegionMap detectVramRegionMap(MemoryDescs const& descs);
+};
+
+/// VMM region metadata exchanged between agents for chunk boundary calculations.
+struct VramRegionMeta
+{
+    uintptr_t baseAddr;
+    size_t totalLen;
+    size_t chunkSize; ///< 0 = cudaMalloc (no split), >0 = VMM chunk size
+};
+
 // `AgentDesc` represents the unique identifier for reading and writing to the agent.
 // By accessing this identifier, the backend can establish the correct connection.
+// It also carries VMM region metadata so that remote agents can split at chunk boundaries.
 class AgentDesc final
 {
 public:
@@ -204,13 +254,31 @@ class AgentDesc final
     {
     }
 
+    AgentDesc(std::string backendAgentDesc, std::vector<VramRegionMeta> vramRegions)
+        : mBackendAgentDesc{std::move(backendAgentDesc)}
+        , mVramRegions{std::move(vramRegions)}
+    {
+    }
+
     [[nodiscard]] std::string const& getBackendAgentDesc() const noexcept
     {
         return mBackendAgentDesc;
     }
 
+    [[nodiscard]] std::vector<VramRegionMeta> const& getVramRegions() const noexcept
+    {
+        return mVramRegions;
+    }
+
+    /// Serialize the entire AgentDesc (backend blob + VMM regions) into an opaque string.
+    [[nodiscard]] std::string serialize() const;
+
+    /// Deserialize an opaque string back into an AgentDesc.
+    [[nodiscard]] static AgentDesc deserialize(std::string const& data);
+
 private:
     std::string mBackendAgentDesc;
+    std::vector<VramRegionMeta> mVramRegions;
 };
 
 // `TransferOp` is an enumeration that represents the types of transfer operations.

cpp/tensorrt_llm/CMakeLists.txt

@@ -200,6 +200,8 @@ set(TRTLLM_LINK_LIBS
     layers_src
     runtime_src
     testing_src
+    mhcKernels_src
+    compressorKernels_src
     userbuffers_src
     ${DECODER_SHARED_TARGET_0}
     ${DECODER_SHARED_TARGET_1})

cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.h

@@ -77,6 +77,18 @@ class IndexMapper
     at::Tensor getCopyIndex(
         std::vector<LlmRequest::RequestIdType> const& requestIds, SizeType32 numContext, SizeType32 beamWidth);
 
+    /// Number of sequences currently tracked (i.e. active IndexMapper slots).
+    [[nodiscard]] SizeType32 size() const noexcept
+    {
+        return static_cast<SizeType32>(indexMap_.size());
+    }
+
+    /// Number of free IndexMapper slots available for new sequences.
+    [[nodiscard]] SizeType32 numFreeSlots() const noexcept
+    {
+        return static_cast<SizeType32>(freeIndices_.size());
+    }
+
 private:
     std::unordered_map<LlmRequest::RequestIdType, SizeType32> indexMap_;
     std::set<SizeType32> freeIndices_;

cpp/tensorrt_llm/common/attentionOp.cpp

@@ -749,7 +749,7 @@ size_t AttentionOp::getFmhaMultiCtasKvScratchSize() const noexcept
 }
 
 size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t max_num_seq, int32_t input_seq_length,
-    int32_t cross_kv_length, int32_t max_num_tokens) const noexcept
+    int32_t cross_kv_length, int32_t max_num_tokens, int32_t total_kv_len) const noexcept
 {
     if (max_num_tokens == 0)
     {
@@ -829,8 +829,12 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t
         }
         else
         {
-            fp8_k_buf_size = mChunkPrefillBufferBatchSize * max_num_tokens * static_cast<size_t>(total_k_dim_all_heads);
-            fp8_v_buf_size = mChunkPrefillBufferBatchSize * max_num_tokens * static_cast<size_t>(total_v_dim_all_heads);
+            // Use total_kv_len when available (KV cache reuse causes total_kv_len >> max_num_tokens).
+            // enqueueContext sizes these buffers by total_kv_len, so workspace must match.
+            size_t const kv_buf_tokens = std::max(
+                static_cast<size_t>(total_kv_len), static_cast<size_t>(mChunkPrefillBufferBatchSize) * max_num_tokens);
+            fp8_k_buf_size = kv_buf_tokens * static_cast<size_t>(total_k_dim_all_heads);
+            fp8_v_buf_size = kv_buf_tokens * static_cast<size_t>(total_v_dim_all_heads);
         }
     }
     else if (useSageAttnSeparateQkv)

cpp/tensorrt_llm/common/attentionOp.h

@@ -57,7 +57,7 @@ class AttentionOp
     [[nodiscard]] int getHeadSize(bool checkInit = true) const;
     [[nodiscard]] int getMaxNumSeqLenTile(int batch_beam_size = 1) const;
     [[nodiscard]] size_t getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t nbReq, int32_t max_input_length,
-        int32_t cross_kv_length = 0, int32_t max_num_tokens = 0) const noexcept;
+        int32_t cross_kv_length = 0, int32_t max_num_tokens = 0, int32_t total_kv_len = 0) const noexcept;
     // total_num_seq is the sum of beam_width for multiple requests
     [[nodiscard]] size_t getWorkspaceSizeForGeneration(nvinfer1::DataType type, int32_t total_num_seq,
         int32_t max_attention_window_size, int32_t max_num_tokens, int32_t max_blocks_per_sequence) const noexcept;

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt

@@ -39,9 +39,8 @@ if(NIXL_ROOT)
   # Link against all NIXL libraries
   target_link_libraries(${NIXL_WRAPPER_TARGET} PRIVATE NIXL::nixl)
 
-  # Link against CUDA
+  # Link against CUDA runtime (for cudaMemcpy in posix fallback)
   target_link_libraries(${NIXL_WRAPPER_TARGET} PRIVATE CUDA::cudart)
-  target_link_libraries(${NIXL_WRAPPER_TARGET} PRIVATE CUDA::cuda_driver)
 
   set(NIXL_ENABLED TRUE)
 else()

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp

@@ -140,7 +140,21 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m)
             {
                 auto const& desc = self.getBackendAgentDesc();
                 return nb::bytes(desc.data(), desc.size());
-            });
+            })
+        .def("serialize",
+            [](kvc::AgentDesc const& self)
+            {
+                auto s = self.serialize();
+                return nb::bytes(s.data(), s.size());
+            })
+        .def_static(
+            "deserialize",
+            [](nb::bytes data)
+            {
+                std::string str(data.c_str(), data.size());
+                return kvc::AgentDesc::deserialize(str);
+            },
+            nb::arg("data"));
 
     // TransferRequest class
     //
@@ -200,39 +214,52 @@ NB_MODULE(tensorrt_llm_transfer_agent_binding, m)
         .def_rw("backend_params", &kvc::BaseAgentConfig::backendParams);
 
     // BaseTransferAgent class (abstract base)
+    // All transfer-engine operations release the GIL: they may block on NIXL /
+    // UCX / network / GPU memory pinning, and holding the GIL across them
+    // starves Python listener / progress threads in the same process.
     nb::class_<kvc::BaseTransferAgent>(m, "BaseTransferAgent")
-        .def("register_memory", &kvc::BaseTransferAgent::registerMemory, nb::arg("descs"))
-        .def("deregister_memory", &kvc::BaseTransferAgent::deregisterMemory, nb::arg("descs"))
+        .def("register_memory", &kvc::BaseTransferAgent::registerMemory, nb::arg("descs"),
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("deregister_memory", &kvc::BaseTransferAgent::deregisterMemory, nb::arg("descs"),
+            nb::call_guard<nb::gil_scoped_release>())
         .def("load_remote_agent",
             nb::overload_cast<std::string const&, kvc::AgentDesc const&>(&kvc::BaseTransferAgent::loadRemoteAgent),
-            nb::arg("name"), nb::arg("agent_desc"))
+            nb::arg("name"), nb::arg("agent_desc"), nb::call_guard<nb::gil_scoped_release>())
         .def("load_remote_agent_by_connection",
             nb::overload_cast<std::string const&, kvc::ConnectionInfoType const&>(
                 &kvc::BaseTransferAgent::loadRemoteAgent),
-            nb::arg("name"), nb::arg("connection_info"))
-        .def("get_local_agent_desc", &kvc::BaseTransferAgent::getLocalAgentDesc)
-        .def("invalidate_remote_agent", &kvc::BaseTransferAgent::invalidateRemoteAgent, nb::arg("name"))
+            nb::arg("name"), nb::arg("connection_info"), nb::call_guard<nb::gil_scoped_release>())
+        .def("get_local_agent_desc", &kvc::BaseTransferAgent::getLocalAgentDesc,
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("invalidate_remote_agent", &kvc::BaseTransferAgent::invalidateRemoteAgent, nb::arg("name"),
+            nb::call_guard<nb::gil_scoped_release>())
         .def(
             "submit_transfer_requests",
             [](kvc::BaseTransferAgent& self, kvc::TransferRequest const& request)
             { return self.submitTransferRequests(request).release(); },
-            nb::arg("request"), nb::rv_policy::take_ownership)
-        .def(
-            "notify_sync_message", &kvc::BaseTransferAgent::notifySyncMessage, nb::arg("name"), nb::arg("sync_message"))
-        .def("get_notified_sync_messages", &kvc::BaseTransferAgent::getNotifiedSyncMessages)
-        .def("get_local_connection_info", &kvc::BaseTransferAgent::getLocalConnectionInfo)
-        .def("check_remote_descs", &kvc::BaseTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs"));
+            nb::arg("request"), nb::rv_policy::take_ownership, nb::call_guard<nb::gil_scoped_release>())
+        .def("notify_sync_message", &kvc::BaseTransferAgent::notifySyncMessage, nb::arg("name"),
+            nb::arg("sync_message"), nb::call_guard<nb::gil_scoped_release>())
+        .def("get_notified_sync_messages", &kvc::BaseTransferAgent::getNotifiedSyncMessages,
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("get_local_connection_info", &kvc::BaseTransferAgent::getLocalConnectionInfo,
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("check_remote_descs", &kvc::BaseTransferAgent::checkRemoteDescs, nb::arg("name"), nb::arg("memory_descs"),
+            nb::call_guard<nb::gil_scoped_release>());
 
 #ifdef ENABLE_NIXL
     // NixlTransferStatus class - release GIL for blocking operations
     nb::class_<kvc::NixlTransferStatus, kvc::TransferStatus>(m, "NixlTransferStatus")
         .def("is_completed", &kvc::NixlTransferStatus::isCompleted, nb::call_guard<nb::gil_scoped_release>())
         .def("wait", &kvc::NixlTransferStatus::wait, nb::arg("timeout_ms") = -1,
-            nb::call_guard<nb::gil_scoped_release>());
+            nb::call_guard<nb::gil_scoped_release>())
+        .def("get_last_status", &kvc::NixlTransferStatus::getLastStatus)
+        .def("get_last_status_str", &kvc::NixlTransferStatus::getLastStatusStr);
 
     // NixlTransferAgent class
     nb::class_<kvc::NixlTransferAgent, kvc::BaseTransferAgent>(m, "NixlTransferAgent")
-        .def(nb::init<kvc::BaseAgentConfig const&>(), nb::arg("config"))
+        .def(nb::init<kvc::BaseAgentConfig const&>(), nb::arg("config"), nb::call_guard<nb::gil_scoped_release>())
+        .def("shutdown", &kvc::NixlTransferAgent::shutdown, nb::call_guard<nb::gil_scoped_release>())
         .def("register_memory", &kvc::NixlTransferAgent::registerMemory, nb::arg("descs"))
         .def("deregister_memory", &kvc::NixlTransferAgent::deregisterMemory, nb::arg("descs"))
         .def("load_remote_agent",