[None][chore] Mass merge release/1.3 branch back to main (batch 2)

cpp/kernels/fmha_v2/setup.py

@@ -2199,8 +2199,16 @@ def selected_mask_types(kspec):
             custom_mask = '0'
         # encoder models (head_size = 32 / 64 / 128) need packed_qkv input layout + padding mask.
         elif kspec.input_layout == InputLayout.PACKED_QKV:
-            # NOTE: 72/80 are added for vision transformer
-            if kspec.head_size not in [32, 64, 72, 80, 128]:
+            # NOTE: 72/80 are added for vision transformers. Pixtral's head size
+            # 104 uses warp specialization on SM90 and non-tiled kernels on SM120.
+            pixtral_padding = kspec.head_size == 104 and (
+                kspec.sm == 100 or
+                (kspec.sm in (90, 120) and kspec.dtype in
+                 ('fp16', 'bf16', 'fp16_fp32') and not kspec.enable_skip_softmax
+                 and kspec.warp_specialization == (kspec.sm == 90) and
+                 (kspec.sm != 120 or not kspec.tiled)))
+            if (kspec.head_size not in [32, 64, 72, 80, 128]
+                    and not pixtral_padding):
                 padding_mask = '0'
         # only cross attention (head_size = 32/64/128) needs contiguous_q_kv input layout + padding mask / custom_mask.
         elif kspec.input_layout == InputLayout.CONTIGUOUS_Q_KV:
@@ -6920,10 +6928,10 @@ def enumerate_kernels():
                   and kspec.cross_mha     == False
                   and kspec.flash_attention == True
                   and kspec.input_layout != InputLayout.SEPARATE_Q_K_V)
-                  # Gemma3 VL support.
+                  # Gemma3 VL and Pixtral VL support.
                   or  (kspec.sm           == 100
                   and kspec.dtype         in ['fp16', 'bf16', 'fp16_fp32', 'e4m3', 'e4m3_fp32']
-                  and kspec.head_size     == 72
+                  and kspec.head_size     in [72, 104]
                   and kspec.head_size_v   == 0
                   and kspec.sage_block_sizes is None
                   and kspec.version       == 2

cpp/tensorrt_llm/common/ncclUtils.cpp

@@ -24,6 +24,7 @@
 #include <cstdlib>
 #include <limits>
 #include <stdexcept>
+#include <utility>
 
 namespace
 {
@@ -398,6 +399,7 @@ size_t NcclCommResourceManager::getResourceCount(ncclComm_t comm) const noexcept
 namespace
 {
 thread_local std::optional<uint64_t> gGraphPoolOwner;
+thread_local bool gWindowAllocationSuppressed = false;
 } // namespace
 
 NCCLWindowAllocator& NCCLWindowAllocator::getInstance()
@@ -417,6 +419,11 @@ void NCCLWindowAllocator::setGraphPoolOwner(int64_t owner)
     gGraphPoolOwner = owner < 0 ? std::nullopt : std::optional<uint64_t>{static_cast<uint64_t>(owner)};
 }
 
+bool NCCLWindowAllocator::setWindowAllocationSuppressed(bool suppressed)
+{
+    return std::exchange(gWindowAllocationSuppressed, suppressed);
+}
+
 void NCCLWindowAllocator::releaseGraphPoolOwner(int64_t owner)
 {
     TLLM_CHECK_WITH_INFO(owner >= 0, "NCCL window graph-pool owner must be non-negative");
@@ -469,6 +476,12 @@ NCCLWindowBuffer NCCLWindowAllocator::requestBuffer(ncclComm_t comm, size_t size
     {
         return NCCLWindowBuffer();
     }
+    if (gWindowAllocationSuppressed)
+    {
+        TLLM_LOG_TRACE("[NCCLUtil] NCCL window allocation suppressed on this thread for comm %p (requested: %zu)",
+            static_cast<void*>(comm), size);
+        return NCCLWindowBuffer();
+    }
 
     TLLM_CHECK_WITH_INFO(comm != nullptr, "NCCL communicator cannot be null");
     TLLM_CHECK_WITH_INFO(size > 0, "Buffer size must be greater than 0");

cpp/tensorrt_llm/common/ncclUtils.h

@@ -247,6 +247,11 @@ class NCCLWindowAllocator
     // still in use are returned when their deterministic tensor scope ends.
     void releaseGraphPoolOwner(int64_t owner);
 
+    // While set on this host thread, requestBuffer returns an invalid buffer without touching the
+    // pool, so callers fall back to plain memory. Used where ranks are not SPMD, such as PARALLEL
+    // autotuning. Returns the previous setting.
+    bool setWindowAllocationSuppressed(bool suppressed);
+
     // Search for a buffer by pointer. Returns an invalid buffer if not found.
     // This matches the UBManager.search_buffer() interface.
     NCCLWindowBuffer searchBuffer(ncclComm_t comm, void* ptr) const;

cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp

@@ -1234,7 +1234,11 @@ bool NixlTransferAgent::checkRemoteDescs(std::string const& name, MemoryDescs co
 {
     std::shared_lock<std::shared_mutex> lock(mLock);
     TLLM_CHECK_WITH_INFO(!mShutdown.load(), "NixlTransferAgent::checkRemoteDescs called after shutdown");
-    auto status = mRawAgent->checkRemoteMD(name, NixlHelper::convertXferDist(memoryDescs));
+    static VramRegionMap const kEmptyMap;
+    auto const remoteIt = mRemoteVramRegionInfo.find(name);
+    auto const& remoteRegionMap = (remoteIt != mRemoteVramRegionInfo.end()) ? remoteIt->second : kEmptyMap;
+    auto const splitDescs = VmmDescSplitter::splitDescsWithRegionMap(memoryDescs, remoteRegionMap);
+    auto status = mRawAgent->checkRemoteMD(name, NixlHelper::convertXferDist(splitDescs));
     TLLM_CHECK_WITH_INFO(status == NIXL_SUCCESS || status == NIXL_ERR_NOT_FOUND, "checkRemoteMD failed with status: %s",
         nixlEnumStrings::statusStr(status).c_str());
     return status == NIXL_SUCCESS;

cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp

@@ -63,10 +63,16 @@ QkvLayout AttentionInputLayoutToQkvLayout(AttentionInputLayout layout)
 
 FmhaDispatcher::FmhaDispatcher(MHARunnerFixedParams fixedParams)
     : mFixedParams(fixedParams)
-    // TRTLLM-GEN only supports power of 2 head sizes (and 80 with padding).
-    // The exception will fall back to fmha v2.
+    // FMHA v2 provides head-size 72 and FP16/BF16 full/causal packed-QKV head-size 104 kernels.
+    // Keep mixed Q/KV precisions on TRTLLM-GEN: fmha v2 requires matching types.
     // Please update fmha_v2/setup.py if you want to add more supported head sizes.
-    , mUseTllmGen(tensorrt_llm::common::isSM100Family() && fixedParams.headSize != 72)
+    , mUseTllmGen(tensorrt_llm::common::isSM100Family() && fixedParams.headSize != 72
+          && !(fixedParams.headSize == 104 && fixedParams.headSizeV == 104
+              && fixedParams.attentionInputLayout == AttentionInputLayout::PACKED_QKV
+              && fixedParams.dataType == fixedParams.dataTypeKv
+              && (fixedParams.dataType == DATA_TYPE_FP16 || fixedParams.dataType == DATA_TYPE_BF16)
+              && (fixedParams.attentionMaskType == ContextAttentionMaskType::PADDING
+                  || fixedParams.attentionMaskType == ContextAttentionMaskType::CAUSAL)))
     , mMultiProcessorCount(tensorrt_llm::common::getMultiProcessorCount())
 {
     if (mUseTllmGen)

cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, 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.
@@ -166,20 +166,37 @@ void NumaHugePagePoolAllocator::shutdown()
     munmap(mMmapBasePtr, kReservedVirtualMemorySize + kHugePageSize);
 }
 
+namespace
+{
+
+//! Whether the host can access device-resident managed memory directly, without page faults or migration.
+//! This holds on hardware-coherent platforms such as Grace-based GB200/GB300, including when GPU memory is not
+//! exposed as a NUMA node (CDMM mode).
+bool isHostDirectManagedAccessSupported()
+{
+    int deviceId = -1;
+    TLLM_CUDA_CHECK(cudaGetDevice(&deviceId));
+    int directAccess = 0;
+    TLLM_CUDA_CHECK(cudaDeviceGetAttribute(&directAccess, cudaDevAttrDirectManagedMemAccessFromHost, deviceId));
+    return directAccess != 0;
+}
+
+} // namespace
+
 bool HostAccessibleDeviceAllocator::mAllowManagedFallback = false;
 
 bool HostAccessibleDeviceAllocator::isSupported()
 {
-    if (!tensorrt_llm::common::getEnvEplbForceGdrcopy()
-        && TopologyDetector::getInstance().getCurrentGpuMemoryNumaId() >= 0)
+    bool const forceGdrcopy = tensorrt_llm::common::getEnvEplbForceGdrcopy();
+    if (!forceGdrcopy && TopologyDetector::getInstance().getCurrentGpuMemoryNumaId() >= 0)
     {
         // we are on systems that GPU memory is also a NUMA node.
         return true;
     }
     if (!tensorrt_llm::runtime::gdrcopy::isInitialized() && !tensorrt_llm::runtime::gdrcopy::initialize())
     {
-        // system don't support GDRCopy.
-        return mAllowManagedFallback;
+        // System doesn't support GDRCopy, managed memory still works if the host can access it directly.
+        return mAllowManagedFallback || (!forceGdrcopy && isHostDirectManagedAccessSupported());
     }
     return true;
 }
@@ -207,6 +224,7 @@ void HostAccessibleDeviceAllocator::init()
         mGpuMemNumaId = TopologyDetector::getInstance().getCurrentGpuMemoryNumaId();
     }
 
+    mHostDirectManagedAccess = false;
     if (mGpuMemNumaId < 0)
     {
         // We only use GDRCopy when there is no NUMA node for GPU memory.
@@ -219,6 +237,13 @@ void HostAccessibleDeviceAllocator::init()
         {
             mGdrHandle = tensorrt_llm::runtime::gdrcopy::open();
         }
+        else if (!tensorrt_llm::common::getEnvEplbForceGdrcopy() && isHostDirectManagedAccessSupported())
+        {
+            mHostDirectManagedAccess = true;
+            TLLM_LOG_WARNING(
+                "GPU memory is not a NUMA node and GDRCopy is not available, using managed memory with direct host "
+                "access as host accessible device memory.");
+        }
     }
     mIsInited = true;
 }
@@ -372,11 +397,17 @@ void* HostAccessibleDeviceAllocator::allocate(size_t memorySize)
     }
     else
     {
-        TLLM_CHECK_WITH_INFO(
-            mAllowManagedFallback, "HostAccessibleDeviceAllocator is not supported on the current system.");
+        TLLM_CHECK_WITH_INFO(mAllowManagedFallback || mHostDirectManagedAccess,
+            "HostAccessibleDeviceAllocator is not supported on the current system.");
         TLLM_CUDA_CHECK(cudaMallocManaged(&devPtr, memorySize));
         TLLM_CUDA_CHECK(cudaMemAdvise(devPtr, memorySize, cudaMemAdviseSetPreferredLocation,
             cudaMemLocation{cudaMemLocationTypeDevice, {currentDevId}}));
+        if (mHostDirectManagedAccess)
+        {
+            // Keep the device-resident pages mapped for the host, so host accesses neither fault nor migrate them.
+            TLLM_CUDA_CHECK(cudaMemAdvise(
+                devPtr, memorySize, cudaMemAdviseSetAccessedBy, cudaMemLocation{cudaMemLocationTypeHost, {0}}));
+        }
         hostPtr = devPtr;
     }
     recordAllocation(devPtr, memorySize, hostPtr, memDesc);
@@ -401,8 +432,8 @@ void HostAccessibleDeviceAllocator::free(void* ptr)
         }
         else
         {
-            TLLM_CHECK_WITH_INFO(
-                mAllowManagedFallback, "HostAccessibleDeviceAllocator is not supported on the current system.");
+            TLLM_CHECK_WITH_INFO(mAllowManagedFallback || mHostDirectManagedAccess,
+                "HostAccessibleDeviceAllocator is not supported on the current system.");
             TLLM_CUDA_CHECK(cudaFree(ptr));
         }
         void* hostPtr = it->second.hostPtr;

cpp/tensorrt_llm/runtime/moeLoadBalancer/hostAccessibleDeviceAllocator.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2024, 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.
@@ -157,10 +157,13 @@ class HostAccessibleDeviceAllocator
      */
     AllocationInfo getAllocationInfoFromDevPtr(void const* devPtr);
 
-    // if GPU memory has NUMA id, then CPU can direct access that. We should use this.
+    // Host access to device memory, in order of preference:
+    // 1. GPU memory is a NUMA node (mGpuMemNumaId >= 0), so the CPU can access it directly.
     int mGpuMemNumaId = -1;
-    // if Not, we should use GDRCopy
+    // 2. Otherwise, GDRCopy maps device memory for the host.
     gdr_t mGdrHandle = nullptr;
+    // 3. Without GDRCopy, managed memory that the host can access directly (hardware-coherent platforms).
+    bool mHostDirectManagedAccess = false;
 
     int mDevId = -1;
 

cpp/tensorrt_llm/thop/allreduceOp.cpp

@@ -2609,6 +2609,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
         "end_nccl_window_tensor_scope.tracked(Tensor[] inputs, Tensor[] outputs, bool failed, int entry_depth) -> ()");
     m.def("set_nccl_window_graph_owner(int owner) -> ()");
     m.def("release_nccl_window_graph_owner(int owner) -> ()");
+    m.def("set_nccl_window_allocation_suppressed(bool suppressed) -> bool");
     m.def("get_nccl_window_buffer_release_epoch() -> int");
     m.def("promote_nccl_window_buffer_releases(int device, int release_epoch) -> ()");
     m.def(
@@ -2671,6 +2672,12 @@ TORCH_LIBRARY_IMPL(trtllm, CatchAll, m)
     m.impl("release_nccl_window_graph_owner",
         [](int64_t owner)
         { tensorrt_llm::common::nccl_util::NCCLWindowAllocator::getInstance().releaseGraphPoolOwner(owner); });
+    m.impl("set_nccl_window_allocation_suppressed",
+        [](bool suppressed)
+        {
+            return tensorrt_llm::common::nccl_util::NCCLWindowAllocator::getInstance().setWindowAllocationSuppressed(
+                suppressed);
+        });
     m.impl("get_nccl_window_buffer_release_epoch",
         []()
         {
@@ -2687,6 +2694,7 @@ TORCH_LIBRARY_IMPL(trtllm, CatchAll, m)
 #else
     m.impl("set_nccl_window_graph_owner", [](int64_t) {});
     m.impl("release_nccl_window_graph_owner", [](int64_t) {});
+    m.impl("set_nccl_window_allocation_suppressed", [](bool) { return false; });
     m.impl("get_nccl_window_buffer_release_epoch", []() { return int64_t{0}; });
     m.impl("promote_nccl_window_buffer_releases", [](int64_t, int64_t) {});
 #endif

cpp/tests/unit_tests/executor/transferAgentTest.cpp

@@ -21,6 +21,7 @@
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
+#include <algorithm>
 #include <atomic>
 #include <filesystem>
 #include <thread>
@@ -125,6 +126,74 @@ TEST_P(TransferAgentTest, Basic)
     xferAgent0->invalidateRemoteAgent(agent1);
 }
 
+TEST_P(TransferAgentTest, CheckRemoteDescriptorRanges)
+{
+    if (backend != "nixl")
+    {
+        GTEST_SKIP() << "NIXL metadata validation";
+    }
+    std::string const agent0{"agent0"}, agent1{"agent1"};
+    BaseAgentConfig config0{agent0, true, false, true}, config1{agent1, true, false, true};
+    auto xferAgent0 = makeTransferAgent(config0);
+    auto xferAgent1 = makeTransferAgent(config1);
+    std::vector<char> source(64, 10);
+    std::vector<char> destination(128, 1);
+    RegisteredHostMemory regSource(MemoryDescs{MemoryType::kDRAM, {MemoryDesc{source}}}, xferAgent0.get());
+    RegisteredHostMemory regDestination(
+        MemoryDescs{MemoryType::kDRAM, {MemoryDesc{destination.data(), source.size(), 0}}}, xferAgent1.get());
+    xferAgent0->loadRemoteAgent(agent1, xferAgent1->getLocalAgentDesc());
+    ASSERT_TRUE(xferAgent0->checkRemoteDescs(agent1, regDestination.getDescs()));
+
+    // The backing allocation is large enough, but the requested range exceeds its registration.
+    MemoryDescs const overrun{MemoryType::kDRAM, {MemoryDesc{destination.data() + 32, source.size(), 0}}};
+    EXPECT_FALSE(xferAgent0->checkRemoteDescs(agent1, overrun));
+    EXPECT_EQ(destination, std::vector<char>(128, 1));
+
+    EXPECT_FALSE(xferAgent0->checkRemoteDescs("unknown-agent", regDestination.getDescs()));
+
+    MemoryDescs const subSource{MemoryType::kDRAM, {MemoryDesc{source.data(), 16, 0}}};
+    MemoryDescs const subDestination{MemoryType::kDRAM, {MemoryDesc{destination.data() + 16, 16, 0}}};
+    EXPECT_TRUE(xferAgent0->checkRemoteDescs(agent1, subDestination));
+    TransferRequest const validWrite{TransferOp::kWRITE, subSource, subDestination, agent1};
+    auto status = xferAgent0->submitTransferRequests(validWrite);
+    ASSERT_EQ(status->wait(), TransferState::kSUCCESS);
+    auto expected = std::vector<char>(128, 1);
+    std::fill(expected.begin() + 16, expected.begin() + 32, 10);
+    EXPECT_EQ(destination, expected);
+    xferAgent0->invalidateRemoteAgent(agent1);
+}
+
+TEST_P(TransferAgentTest, CheckRemoteDescriptorsAcrossVmmChunks)
+{
+    if (backend != "nixl")
+    {
+        GTEST_SKIP() << "NIXL metadata validation";
+    }
+    constexpr size_t kChunkSize = 64;
+    constexpr size_t kTotalSize = 2 * kChunkSize;
+    char* destination = nullptr;
+    TLLM_CUDA_CHECK(cudaMalloc(&destination, kTotalSize));
+    std::unique_ptr<char, decltype(&cudaFree)> allocation{destination, &cudaFree};
+    std::string const agent0{"agent0"}, agent1{"agent1"};
+    BaseAgentConfig config0{agent0, true, false, true}, config1{agent1, true, false, true};
+    auto xferAgent0 = makeTransferAgent(config0);
+    auto xferAgent1 = makeTransferAgent(config1);
+    RegisteredHostMemory registered(
+        MemoryDescs{MemoryType::kVRAM,
+            {MemoryDesc{destination, kChunkSize, 0}, MemoryDesc{destination + kChunkSize, kChunkSize, 0}}},
+        xferAgent1.get());
+    // Advertise a chunked region over two real registrations without requiring a VMM allocator.
+    auto const metadata = xferAgent1->getLocalAgentDesc();
+    AgentDesc const chunked{
+        metadata.getBackendAgentDesc(), {{reinterpret_cast<uintptr_t>(destination), kTotalSize, kChunkSize}}};
+    xferAgent0->loadRemoteAgent(agent1, chunked);
+    EXPECT_TRUE(
+        xferAgent0->checkRemoteDescs(agent1, MemoryDescs{MemoryType::kVRAM, {MemoryDesc{destination, kTotalSize, 0}}}));
+    EXPECT_FALSE(xferAgent0->checkRemoteDescs(
+        agent1, MemoryDescs{MemoryType::kVRAM, {MemoryDesc{destination, kTotalSize + 1, 0}}}));
+    xferAgent0->invalidateRemoteAgent(agent1);
+}
+
 TEST_P(TransferAgentTest, Basic2)
 {
 

cpp/tests/unit_tests/kernels/CMakeLists.txt

@@ -101,6 +101,7 @@ endif()
 add_gtest(sparseKvCacheTest sparseKvCacheTest.cu)
 add_gtest(prepareCustomMaskTest prepareCustomMaskTest.cpp)
 add_gtest(fmhaDsv4EpilogueFusionTest fmhaDsv4EpilogueFusionTest.cpp)
+add_gtest(fmhaHead104Test fmhaHead104Test.cpp)
 
 set(NVFP4_COLD_PAGE_KERNEL_TEST_SRC
     nvfp4ColdPageKernelsTest.cpp