[None][perf] Overlap TRTLLM-Gen FMHA JIT warmup compilation with engine warmup (opt-in)

cpp/tensorrt_llm/thop/CMakeLists.txt

@@ -58,6 +58,7 @@ add_th_common_family(
   mlaPreprocessOp.cpp
   attentionOp.cpp
   fmhaPackMaskOp.cpp
+  trtllmGenFmhaJitWarmupOp.cpp
   deepseekV4BlockTableOp.cpp
   inverseRopeFp8QuantOp.cpp
   kdaDecodeOp.cpp

cpp/tensorrt_llm/thop/trtllmGenFmhaJitWarmupOp.cpp

@@ -0,0 +1,74 @@
+/*
+ * Copyright (c) 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "tensorrt_llm/common/config.h"
+#include "tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h"
+
+#include <torch/library.h>
+
+#include <cstdint>
+
+TRTLLM_NAMESPACE_BEGIN
+
+namespace torch_ext
+{
+
+// Waits for the background TRTLLM-Gen FMHA JIT warmup sweeps and re-runs each
+// finished sweep once on the calling thread, so every kernel the warmup
+// derived is compiled before CUDA graph capture and before the engine reports
+// ready. Returns the number of sweeps verified here. No-op unless
+// TRTLLM_GEN_FMHA_ASYNC_WARMUP=1.
+int64_t trtllmGenFmhaJitWarmupDrainAndVerify()
+{
+    return tensorrt_llm::kernels::TllmGenFmhaKernel::drainAndVerifyAllJITWarmups();
+}
+
+// Number of TRTLLM-Gen FMHA kernel-cache misses (NVRTC compiles) the export
+// library reported in this process so far. A value that does not change across
+// a request proves that the request compiled nothing.
+int64_t trtllmGenFmhaJitNumCacheMisses()
+{
+    return tensorrt_llm::kernels::TllmGenFmhaKernel::numJITCacheMisses();
+}
+
+// Number of compile requests whose cache result the export library did not
+// report. Non-zero means the miss count above is not a complete measure.
+int64_t trtllmGenFmhaJitNumUnknownCacheResults()
+{
+    return tensorrt_llm::kernels::TllmGenFmhaKernel::numJITUnknownCacheResults();
+}
+
+// Whether this process runs the TRTLLM-Gen FMHA JIT warmup on the background
+// thread (TRTLLM_GEN_FMHA_ASYNC_WARMUP=1, read once per process).
+bool trtllmGenFmhaAsyncJitWarmupEnabled()
+{
+    return tensorrt_llm::kernels::asyncJITWarmupEnabled();
+}
+
+} // namespace torch_ext
+
+TRTLLM_NAMESPACE_END
+
+TORCH_LIBRARY_FRAGMENT(trtllm, m)
+{
+    m.def("trtllm_gen_fmha_jit_warmup_drain_and_verify() -> int",
+        &tensorrt_llm::torch_ext::trtllmGenFmhaJitWarmupDrainAndVerify);
+    m.def("trtllm_gen_fmha_jit_num_cache_misses() -> int", &tensorrt_llm::torch_ext::trtllmGenFmhaJitNumCacheMisses);
+    m.def("trtllm_gen_fmha_jit_num_unknown_cache_results() -> int",
+        &tensorrt_llm::torch_ext::trtllmGenFmhaJitNumUnknownCacheResults);
+    m.def("trtllm_gen_fmha_async_jit_warmup_enabled() -> bool",
+        &tensorrt_llm::torch_ext::trtllmGenFmhaAsyncJitWarmupEnabled);
+}

cpp/tests/unit_tests/kernels/CMakeLists.txt

@@ -104,6 +104,9 @@ endif()
 
 add_gtest(sparseKvCacheTest sparseKvCacheTest.cu)
 add_gtest(prepareCustomMaskTest prepareCustomMaskTest.cpp)
+add_gtest(fmhaJitWarmupTest fmhaJitWarmupTest.cpp)
+add_gtest(fmhaJitWarmupFingerprintTest fmhaJitWarmupFingerprintTest.cpp)
+add_gtest(fmhaJitWarmupProcessExitTest fmhaJitWarmupProcessExitTest.cpp)
 add_gtest(fmhaDsv4EpilogueFusionTest fmhaDsv4EpilogueFusionTest.cpp)
 
 set(NVFP4_COLD_PAGE_KERNEL_TEST_SRC

cpp/tests/unit_tests/kernels/fmhaJitWarmupFingerprintTest.cpp

@@ -0,0 +1,145 @@
+/*
+ * Copyright (c) 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Tests for TllmGenFmhaKernel::jitWarmupFingerprint(), the key that
+// deduplicates warmup sweeps (in sync and async mode alike). Two sweeps that
+// would compile different kernels must never share a fingerprint. No kernel is
+// loaded or launched: the kernel objects here carry an empty kernel table.
+
+#include "tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h"
+
+#include <gtest/gtest.h>
+
+#include <memory>
+
+namespace
+{
+
+using tensorrt_llm::kernels::Data_type;
+using tensorrt_llm::kernels::FmhaKernelType;
+using tensorrt_llm::kernels::kSM_100;
+using tensorrt_llm::kernels::TllmGenFmhaKernel;
+
+using RunnerParams = TllmGenFmhaKernel::RunnerParams;
+
+// Kernel objects are not copyable; hand them out on the heap.
+std::unique_ptr<TllmGenFmhaKernel> makeKernel(Data_type dtypeOut = Data_type::DATA_TYPE_BF16)
+{
+    return std::make_unique<TllmGenFmhaKernel>(/*pMetaStart=*/nullptr, /*nMetaCount=*/0, Data_type::DATA_TYPE_BF16,
+        Data_type::DATA_TYPE_BF16, Data_type::DATA_TYPE_BF16, dtypeOut, kSM_100);
+}
+
+// A generation-phase sweep request as the warmup forward issues it.
+RunnerParams makeParams()
+{
+    RunnerParams params{};
+    params.mKernelType = FmhaKernelType::Generation;
+    params.mHeadDimQk = 128;
+    params.mHeadDimV = 128;
+    params.mNumHeadsQ = 32;
+    params.mNumHeadsKv = 8;
+    params.mNumHeadsQPerKv = 4;
+    params.mMaxSeqLenQ = 1;
+    params.mNumTokensPerPage = 64;
+    params.mMaxNumPagesPerSeqKv = 256;
+    params.mJITWarmupMaxNumRequests = 64;
+    params.mJITWarmupMaxSeqLenQ = 1;
+    params.mJITWarmupMaxSeqLenKv = 16384;
+    params.mMultiProcessorCount = 148;
+    return params;
+}
+
+TEST(TllmGenFmhaJitWarmupFingerprint, IsDeterministicForIdenticalRequests)
+{
+    auto const kernel = makeKernel();
+    auto const params = makeParams();
+    EXPECT_EQ(kernel->jitWarmupFingerprint(params, 0), kernel->jitWarmupFingerprint(params, 0));
+    // A second kernel object with the same configuration derives the same key.
+    EXPECT_EQ(makeKernel()->jitWarmupFingerprint(params, 0), kernel->jitWarmupFingerprint(params, 0));
+}
+
+TEST(TllmGenFmhaJitWarmupFingerprint, IgnoresPerCallPointersAndStream)
+{
+    auto const kernel = makeKernel();
+    auto const base = makeParams();
+    auto withPointers = base;
+    withPointers.stream = reinterpret_cast<cudaStream_t>(0x1);
+    withPointers.qPtr = reinterpret_cast<void const*>(0x2);
+    withPointers.kvPtr = reinterpret_cast<void const*>(0x3);
+    EXPECT_EQ(kernel->jitWarmupFingerprint(withPointers, 0), kernel->jitWarmupFingerprint(base, 0));
+}
+
+TEST(TllmGenFmhaJitWarmupFingerprint, DistinguishesDevices)
+{
+    auto const kernel = makeKernel();
+    auto const params = makeParams();
+    EXPECT_NE(kernel->jitWarmupFingerprint(params, 0), kernel->jitWarmupFingerprint(params, 1));
+}
+
+TEST(TllmGenFmhaJitWarmupFingerprint, DistinguishesKernelObjectConfigurations)
+{
+    auto const params = makeParams();
+    EXPECT_NE(makeKernel(Data_type::DATA_TYPE_BF16)->jitWarmupFingerprint(params, 0),
+        makeKernel(Data_type::DATA_TYPE_E4M3)->jitWarmupFingerprint(params, 0));
+}
+
+TEST(TllmGenFmhaJitWarmupFingerprint, DistinguishesKernelSelectingRunnerParams)
+{
+    auto const kernel = makeKernel();
+    auto const base = makeParams();
+    auto const baseKey = kernel->jitWarmupFingerprint(base, 0);
+
+    // Eagle3 one-model: target layers (Q > 1, causal spec-decoding kernel) and
+    // draft layers (Q = 1) must get separate sweeps.
+    auto specTarget = base;
+    specTarget.mMaxSeqLenQ = 5;
+    EXPECT_NE(kernel->jitWarmupFingerprint(specTarget, 0), baseKey);
+
+    auto specTree = specTarget;
+    specTree.mIsSpecDecTree = true;
+    EXPECT_NE(kernel->jitWarmupFingerprint(specTree, 0), kernel->jitWarmupFingerprint(specTarget, 0));
+
+    auto wideGrid = base;
+    wideGrid.mJITWarmupMaxSeqLenKv = base.mJITWarmupMaxSeqLenKv * 2;
+    EXPECT_NE(kernel->jitWarmupFingerprint(wideGrid, 0), baseKey);
+
+    auto moreHeads = base;
+    moreHeads.mNumHeadsQPerKv = 32;
+    EXPECT_NE(kernel->jitWarmupFingerprint(moreHeads, 0), baseKey);
+
+    auto sparse = base;
+    sparse.mSparseTopK = 2048;
+    EXPECT_NE(kernel->jitWarmupFingerprint(sparse, 0), baseKey);
+}
+
+TEST(TllmGenFmhaJitWarmupFingerprint, DistinguishesDsv4EpilogueFusion)
+{
+    // The fused epilogue selects a different kernel (and forces an E4M3
+    // output), so a fused and a non-fused sweep with otherwise identical
+    // params, on an E4M3-output kernel object, must not deduplicate to one.
+    auto const kernel = makeKernel(Data_type::DATA_TYPE_E4M3);
+    auto const unfused = makeParams();
+    auto fused = unfused;
+    fused.mDsv4EpilogueFusion.enabled = true;
+    EXPECT_NE(kernel->jitWarmupFingerprint(fused, 0), kernel->jitWarmupFingerprint(unfused, 0));
+
+    // The fusion's per-call buffers do not take part.
+    auto fusedWithBuffers = fused;
+    fusedWithBuffers.mDsv4EpilogueFusion.cosSinCache = reinterpret_cast<float const*>(0x4);
+    EXPECT_EQ(kernel->jitWarmupFingerprint(fusedWithBuffers, 0), kernel->jitWarmupFingerprint(fused, 0));
+}
+
+} // namespace

cpp/tests/unit_tests/kernels/fmhaJitWarmupProcessExitTest.cpp

@@ -0,0 +1,91 @@
+/*
+ * Copyright (c) 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Process-exit test for the asynchronous TRTLLM-Gen FMHA JIT warmup worker, run
+// in a child process (gtest death test). The worker is created first, through
+// the pre-capture barrier with nothing queued, before the install-path string
+// or any other sweep dependency has been touched; work is queued afterwards
+// and is still pending when the process exits. The child exits 0 only if that
+// pending work ran to completion during exit and observed the same install path
+// the main thread saw, i.e. the worker was shut down with its dependencies
+// intact regardless of which path created it.
+
+#include "tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h"
+
+#include <gtest/gtest.h>
+
+#include <atomic>
+#include <chrono>
+#include <cstdlib>
+#include <string>
+#include <thread>
+
+namespace
+{
+
+using tensorrt_llm::kernels::asyncJITWarmupWorkerCreated;
+using tensorrt_llm::kernels::getAsyncJITWarmupWorker;
+using tensorrt_llm::kernels::TllmGenFmhaKernel;
+
+std::atomic<bool> gTaskFinished{false};
+std::atomic<bool> gTaskSawSameExecPath{false};
+std::string gMainThreadExecPath;
+
+[[noreturn]] void runChild()
+{
+    setenv("TRTLLM_GEN_FMHA_ASYNC_WARMUP", "1", 1);
+
+    // Registered first, so it runs last, after the worker's own atexit shutdown.
+    std::atexit(
+        []()
+        {
+            bool const ok = gTaskFinished.load() && gTaskSawSameExecPath.load();
+            std::_Exit(ok ? 0 : 42);
+        });
+
+    // "Empty barrier": nothing has been queued and the install path has never
+    // been computed. This must not create a worker whose shutdown ordering
+    // depends on static objects that do not exist yet.
+    EXPECT_EQ(TllmGenFmhaKernel::drainAndVerifyAllJITWarmups(), 0);
+    EXPECT_FALSE(asyncJITWarmupWorkerCreated().load());
+
+    // Create the worker explicitly before the install path exists, which is the
+    // worst-case creation order, then compute the path on this thread.
+    auto& worker = getAsyncJITWarmupWorker();
+    EXPECT_TRUE(asyncJITWarmupWorkerCreated().load());
+    gMainThreadExecPath = TllmGenFmhaKernel::getExecPath();
+
+    // Queue work that is still pending when exit() starts.
+    bool const queued = worker.enqueue(
+        []()
+        {
+            std::this_thread::sleep_for(std::chrono::milliseconds(300));
+            gTaskSawSameExecPath = (TllmGenFmhaKernel::getExecPath() == gMainThreadExecPath);
+            gTaskFinished = true;
+        });
+    EXPECT_TRUE(queued);
+    EXPECT_EQ(worker.numPendingTasks(), 1);
+
+    // Normal process exit: atexit handlers and static destruction run.
+    std::exit(0);
+}
+
+TEST(TllmGenFmhaAsyncJITWarmupProcessExit, PendingWorkFinishesWithDependenciesAliveAtExit)
+{
+    EXPECT_EXIT(runChild(), ::testing::ExitedWithCode(0), "");
+}
+
+} // namespace