diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h new file mode 100644 index 000000000000..1ae3691660c8 --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h @@ -0,0 +1,459 @@ +/* + * 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. + */ + +#pragma once + +#include "tensorrt_llm/common/config.h" +#include "tensorrt_llm/common/logger.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Host-only building blocks for the asynchronous TRTLLM-Gen FMHA JIT warmup: +// a single-thread task worker with an explicitly ordered shutdown, and the +// registry that tracks each warmup sweep from scheduling through foreground +// verification. Neither depends on CUDA, so both are covered by host unit +// tests; the CUDA-specific parts (relaxed stream-capture mode on the worker +// thread, device selection, the sweeps themselves) live in fmhaKernels.h. + +TRTLLM_NAMESPACE_BEGIN + +namespace kernels +{ + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// A single background worker thread. A single thread is intentional: the +// warmup compiles serialize on the FMHA export library's interface mutex +// anyway (its kernel cache is not thread-safe), so a pool would add no +// parallelism. +// +// Lifetime: the process-wide instance is leaked and shut down (drained, then +// joined) through a std::atexit handler registered when the worker is first +// created, so no sweep is still running during static destruction. The +// ordering of that handler against other static objects is not relied upon: every +// object a sweep reads is either leaked (the sweep registry, the export +// library's interface and selection mutexes, the install-path string) or +// constructed when the library is loaded (the export library's own globals), +// and each kernel object drains the worker in its destructor before a sweep +// that captured it could outlive it. Whichever path creates the worker first, +// and however late a sweep is queued, a pending sweep therefore finishes with +// everything it touches alive. +class TllmGenFmhaAsyncWarmupWorker +{ +public: + using Task = std::function; + + // threadInit runs once on the worker thread before any task (e.g. to set + // the CUDA stream-capture mode for that thread). + explicit TllmGenFmhaAsyncWarmupWorker(std::function threadInit = {}) + : mThreadInit(std::move(threadInit)) + { + } + + ~TllmGenFmhaAsyncWarmupWorker() + { + shutdown(); + } + + TllmGenFmhaAsyncWarmupWorker(TllmGenFmhaAsyncWarmupWorker const&) = delete; + TllmGenFmhaAsyncWarmupWorker& operator=(TllmGenFmhaAsyncWarmupWorker const&) = delete; + + // Queues a task for the worker thread. Returns false if the task was NOT + // queued, in which case the caller must run it inline: either the worker + // thread could not be started, or the worker has already been shut down. + bool enqueue(Task task) + { + { + std::lock_guard lock(mMutex); + if (mStopRequested) + { + return false; + } + if (!mThreadStarted) + { + // Only mark started after the thread object is successfully constructed: + // if std::thread throws (resource exhaustion), the task must not be + // stranded behind a worker that never runs (drain() would deadlock). + try + { + mThread = std::thread([this]() { workerLoop(); }); + mThreadStarted = true; + } + catch (std::exception const& e) + { + TLLM_LOG_WARNING( + "TRTLLM-Gen FMHA async JIT warmup: worker thread failed to start (%s); compiling inline.", + e.what()); + return false; + } + } + mQueue.push_back(std::move(task)); + mNumPendingTasks.fetch_add(1, std::memory_order_acq_rel); + } + mQueueCv.notify_one(); + return true; + } + + int numPendingTasks() const + { + return mNumPendingTasks.load(std::memory_order_acquire); + } + + bool isShutDown() const + { + std::lock_guard lock(mMutex); + return mStopRequested; + } + + // Blocks until every queued task has finished. Returns immediately when + // nothing is pending, including after shutdown(). + void drain() + { + std::unique_lock lock(mMutex); + mDrainCv.wait(lock, [this]() { return mNumPendingTasks.load(std::memory_order_acquire) == 0; }); + } + + // Runs every queued task to completion, then stops and joins the worker + // thread. Idempotent. Later enqueue() calls return false so callers fall + // back to inline execution. + void shutdown() + { + std::thread thread; + { + std::unique_lock lock(mMutex); + if (mStopRequested) + { + return; + } + mStopRequested = true; + thread = std::move(mThread); + } + mQueueCv.notify_all(); + if (thread.joinable()) + { + thread.join(); + } + } + +private: + void workerLoop() + { + if (mThreadInit) + { + mThreadInit(); + } + while (true) + { + Task task; + { + std::unique_lock lock(mMutex); + mQueueCv.wait(lock, [this]() { return !mQueue.empty() || mStopRequested; }); + if (mQueue.empty()) + { + return; + } + task = std::move(mQueue.front()); + mQueue.pop_front(); + } + try + { + task(); + } + catch (std::exception const& e) + { + TLLM_LOG_WARNING("TRTLLM-Gen FMHA async JIT warmup task failed: %s", e.what()); + } + catch (...) + { + TLLM_LOG_WARNING("TRTLLM-Gen FMHA async JIT warmup task failed with an unknown error."); + } + { + std::lock_guard lock(mMutex); + mNumPendingTasks.fetch_sub(1, std::memory_order_acq_rel); + } + mDrainCv.notify_all(); + } + } + + std::function mThreadInit; + mutable std::mutex mMutex; + std::condition_variable mQueueCv; + std::condition_variable mDrainCv; + std::deque mQueue; + std::atomic mNumPendingTasks{0}; + std::thread mThread; + bool mThreadStarted{false}; + bool mStopRequested{false}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Per-sweep warmup progress. +enum class JITWarmupState : int +{ + kScheduled = 0, // Background sweep queued. + kBackgroundDone, // Background sweep finished (or failed); foreground verify pass pending. + kVerified, // Sweep ran, or re-ran, on a foreground thread. +}; + +// Tracks every distinct warmup sweep requested in the process, keyed by a +// fingerprint of the sweep parameters. The registry is shared by every kernel +// object on every device, so each entry records the device whose context the +// compiled modules were loaded into and the kernel object that owns the +// sweep: a verify pass must re-run a sweep through the same kernel object, on +// the same device, that scheduled it. +// +// TParams is the sweep parameter struct (copied by value) and TInstance the +// owning kernel type; the registry only stores a pointer to the latter. +template +class TllmGenFmhaJitWarmupRegistry +{ +public: + struct PendingSweep + { + uint64_t mFingerprint; + TParams mParams; + int mDeviceId; + TInstance* mInstance; + }; + + // Registers a sweep. Returns true if an identical sweep is already + // registered (scheduled, finished, or verified), in which case the caller + // must not run it again. + bool tryRegister( + uint64_t fingerprint, TParams const& params, JITWarmupState initialState, int deviceId, TInstance* instance) + { + std::lock_guard lock(mMutex); + if (mEntries.find(fingerprint) != mEntries.end()) + { + return true; + } + mEntries.emplace(fingerprint, Entry{params, initialState, deviceId, instance}); + return false; + } + + // Moves a scheduled sweep to kBackgroundDone. Called when the background + // task finishes, whether or not it succeeded, so a failed sweep is re-run + // by the verify pass instead of being deduplicated away. + void markBackgroundDone(uint64_t fingerprint) + { + std::lock_guard lock(mMutex); + auto it = mEntries.find(fingerprint); + if (it != mEntries.end() && it->second.mState == JITWarmupState::kScheduled) + { + it->second.mState = JITWarmupState::kBackgroundDone; + mNumUnverified.fetch_add(1, std::memory_order_acq_rel); + } + } + + // Number of sweeps in kBackgroundDone state (cheap gate for the verify pass). + int numUnverified() const + { + return mNumUnverified.load(std::memory_order_acquire); + } + + // Claims one kBackgroundDone sweep owned by (deviceId, instance), marking + // it verified. The caller runs the returned sweep. + std::optional claimUnverified(int deviceId, TInstance const* instance) + { + return claim([deviceId, instance](Entry const& entry) + { return entry.mDeviceId == deviceId && entry.mInstance == instance; }); + } + + // Claims one kBackgroundDone sweep regardless of owner, for a caller that + // can switch devices and dispatch to the owning instance (the pre-capture + // drain-and-verify barrier). + std::optional claimAnyUnverified() + { + return claim([](Entry const&) { return true; }); + } + + std::optional state(uint64_t fingerprint) const + { + std::lock_guard lock(mMutex); + auto it = mEntries.find(fingerprint); + if (it == mEntries.end()) + { + return std::nullopt; + } + return it->second.mState; + } + + size_t size() const + { + std::lock_guard lock(mMutex); + return mEntries.size(); + } + +private: + struct Entry + { + TParams mParams; + JITWarmupState mState; + int mDeviceId; + TInstance* mInstance; + }; + + template + std::optional claim(Pred const& pred) + { + std::lock_guard lock(mMutex); + for (auto& [fingerprint, entry] : mEntries) + { + if (entry.mState == JITWarmupState::kBackgroundDone && pred(entry)) + { + entry.mState = JITWarmupState::kVerified; + mNumUnverified.fetch_sub(1, std::memory_order_acq_rel); + return PendingSweep{fingerprint, entry.mParams, entry.mDeviceId, entry.mInstance}; + } + } + return std::nullopt; + } + + mutable std::mutex mMutex; + std::unordered_map mEntries; + std::atomic mNumUnverified{0}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Outcome of scheduleJITWarmupSweep(). +enum class JITWarmupScheduleResult : int +{ + kDuplicate = 0, // An identical sweep was already registered; nothing ran. + kQueued, // Queued on the background worker. + kRanInline, // Ran on the calling thread: async disabled, or the worker could not accept work. +}; + +// Registers one warmup sweep and either queues it on the background worker or +// runs it inline. This is the single scheduling path, shared by the kernel and +// by the protocol tests. +// +// - Deduplication: a sweep whose fingerprint is already registered, whether +// scheduled, finished, or verified, is dropped. This applies in async and in +// sync mode alike, so the fingerprint must cover every input that changes +// the set of kernels a sweep compiles. +// - Async mode (`worker` non-null): the sweep is registered as kScheduled and +// queued. The background task marks it kBackgroundDone whether or not +// `runSweep` throws, so a failed sweep is re-run by the verify pass instead +// of being deduplicated away. If the worker cannot accept work (its thread +// failed to start, or it is shut down), the sweep runs inline and is marked +// done from the calling thread. +// - Sync mode (`worker` null): the sweep is registered as kVerified and run +// inline. +// +// `runSweep(params, deviceId)` performs the sweep; when it runs on the worker +// thread it is responsible for making `deviceId` current. `registry` must +// outlive every queued task. +template +JITWarmupScheduleResult scheduleJITWarmupSweep(TRegistry& registry, TllmGenFmhaAsyncWarmupWorker* worker, + uint64_t fingerprint, TParams const& params, int deviceId, TInstance* instance, TRunSweep runSweep) +{ + bool const async = worker != nullptr; + if (registry.tryRegister( + fingerprint, params, async ? JITWarmupState::kScheduled : JITWarmupState::kVerified, deviceId, instance)) + { + return JITWarmupScheduleResult::kDuplicate; + } + if (!async) + { + runSweep(params, deviceId); + return JITWarmupScheduleResult::kRanInline; + } + bool const queued = worker->enqueue( + [®istry, fingerprint, params, deviceId, runSweep]() + { + try + { + runSweep(params, deviceId); + } + catch (...) + { + registry.markBackgroundDone(fingerprint); + throw; + } + registry.markBackgroundDone(fingerprint); + }); + if (queued) + { + return JITWarmupScheduleResult::kQueued; + } + runSweep(params, deviceId); + registry.markBackgroundDone(fingerprint); + return JITWarmupScheduleResult::kRanInline; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Process-wide tally of the export library's kernel-cache results, one entry +// per generateAndCompileKernel() call. A cache miss is a compile, so a miss +// count that does not change across a request proves the request compiled +// nothing; the same key compiled again after an LRU eviction is a second +// miss. An Unknown result means the library did not report, which a test +// must treat as a failure to measure rather than as "no compile". +class TllmGenFmhaJitCompileStats +{ +public: + enum class CacheResult : int + { + kUnknown = 0, + kHit, + kMiss, + }; + + void record(CacheResult result) + { + switch (result) + { + case CacheResult::kHit: mNumHits.fetch_add(1, std::memory_order_relaxed); break; + case CacheResult::kMiss: mNumMisses.fetch_add(1, std::memory_order_relaxed); break; + case CacheResult::kUnknown: mNumUnknown.fetch_add(1, std::memory_order_relaxed); break; + } + } + + int64_t numMisses() const + { + return mNumMisses.load(std::memory_order_relaxed); + } + + int64_t numHits() const + { + return mNumHits.load(std::memory_order_relaxed); + } + + int64_t numUnknown() const + { + return mNumUnknown.load(std::memory_order_relaxed); + } + +private: + std::atomic mNumMisses{0}; + std::atomic mNumHits{0}; + std::atomic mNumUnknown{0}; +}; + +} // namespace kernels + +TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 79cc1eb5dd7e..b43d211021dc 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -19,18 +19,27 @@ #include "cuda_runtime_api.h" #include "tensorrt_llm/common/config.h" #include +#include #include +#include +#include #include +#include #include +#include +#include #include #include #include #include +#include #include #include +#include #include #include #include +#include #include #include "tensorrt_llm/common/cudaDriverWrapper.h" @@ -38,6 +47,7 @@ #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/kernels/multiHeadAttentionCommon.h" +#include "tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h" #include "cubin/kernelMetaInfo.h" #include "fmhaReduction.h" @@ -103,6 +113,100 @@ constexpr bool isSMCompatible(int gpuSM, int kernelSM) return gpuSM == kernelSM; } +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Serializes every call into the prebuilt trtllm-gen FMHA export library +// (libTrtLlmGenFmhaLib.a). The library caches NVRTC-compiled kernels in an +// unsynchronized, file-scope global LRU cache (fmha::kernelCache); even cache +// hits mutate the LRU list, so FmhaInterface::generateAndCompileKernel() and +// FmhaInterface::run() must never execute concurrently. Every call site in +// this file must hold this mutex. Intentionally leaked: a background warmup +// sweep may still take it during process teardown (see fmhaJitWarmup.h). +inline std::mutex& getFmhaInterfaceMutex() +{ + static auto* const sMutex = new std::mutex(); + return *sMutex; +} + +// Serializes kernel selection (FmhaAutoTuner::selectKernel and +// fmha::computeNumCtas). The selection path queries CUDA occupancy through a +// shared dummy kernel (fmha::emulatedFmhaKernel in FmhaAutoTuner.cu.o) whose +// function attribute it mutates first (cudaFuncSetAttribute followed by +// cudaOccupancyMaxPotentialClusterSize / cudaOccupancyMaxActiveClusters in +// FmhaAutoTuner::getMaxNumActiveClusters). Two threads interleaving +// set-attribute + occupancy-query on that shared handle can read each other's +// smem sizes and select different kernel variants (headDimPerCtaV / +// CGA-reduction / tile decisions) for identical inputs — observed as the +// background warmup sweep compiling a different KernelKey set than foreground +// traffic derives. Selection is microseconds-fast; this lock is never held +// while getFmhaInterfaceMutex() is taken. Intentionally leaked, like the +// interface mutex. +inline std::mutex& getFmhaSelectionMutex() +{ + static auto* const sMutex = new std::mutex(); + return *sMutex; +} + +// Whether the TRTLLM-Gen FMHA JIT warmup grid is compiled on a background +// thread (TRTLLM_GEN_FMHA_ASYNC_WARMUP=1) instead of synchronously inside the +// first warmup forward pass. +inline bool asyncJITWarmupEnabled() +{ + static bool const sEnabled = []() + { + char const* env = std::getenv("TRTLLM_GEN_FMHA_ASYNC_WARMUP"); + return env != nullptr && env[0] == '1' && env[1] == '\0'; + }(); + return sEnabled; +} + +// Whether getAsyncJITWarmupWorker() has created the process-wide worker. Lets +// drain paths (the pre-capture barrier, kernel destruction) skip the worker +// entirely when no sweep was ever queued, instead of creating one just to +// find it idle. +inline std::atomic& asyncJITWarmupWorkerCreated() +{ + static std::atomic sCreated{false}; + return sCreated; +} + +// Process-wide worker for the background JIT warmup sweeps. See +// fmhaJitWarmup.h for the lifetime and shutdown-ordering contract. Every +// dependency a queued sweep reads is either leaked (the registry, the +// interface and selection mutexes, the install path) or constructed at library +// load (the export library's globals), and kernel objects drain the worker in +// their destructor, so it does not matter which path creates the worker first. +inline TllmGenFmhaAsyncWarmupWorker& getAsyncJITWarmupWorker() +{ + static TllmGenFmhaAsyncWarmupWorker* const sWorker = []() + { + auto* worker = new TllmGenFmhaAsyncWarmupWorker( + []() + { + // Relaxed capture mode: CUDA calls made by the worker thread (NVRTC module + // loads) must neither fail nor invalidate a CUDA graph capture that the + // forward-pass thread may have in progress. Module loads never touch a + // capturing stream. This is the same pattern NCCL/cuDNN use for lazy + // module loading from helper threads. + cudaStreamCaptureMode mode = cudaStreamCaptureModeRelaxed; + if (cudaThreadExchangeStreamCaptureMode(&mode) != cudaSuccess) + { + TLLM_LOG_WARNING( + "TRTLLM-Gen FMHA async JIT warmup: failed to set relaxed stream-capture mode; " + "run() will drain this worker before any CUDA graph capture."); + } + }); + // Drains and joins the worker at process exit so no sweep is still running + // during static destruction. The object itself is intentionally leaked. + std::atexit([]() { getAsyncJITWarmupWorker().shutdown(); }); + asyncJITWarmupWorkerCreated().store(true, std::memory_order_release); + return worker; + }(); + return *sWorker; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + inline void checkTrtllmSkipCorrThreshold(fmha::FmhaOptions const& options, float skipCorrThreshold) { fmha::checkSkipCorrThreshold(options, skipCorrThreshold); @@ -175,6 +279,21 @@ class TllmGenFmhaKernel { } + ~TllmGenFmhaKernel() + { + // A queued background warmup task captures `this`, so no such task may + // outlive this object: wait for it here, whether this object dies while + // the process keeps running or during static destruction (the worker and + // everything a sweep reads are leaked, so waiting is safe either way). + if (asyncJITWarmupEnabled() && asyncJITWarmupWorkerCreated().load(std::memory_order_acquire)) + { + getAsyncJITWarmupWorker().drain(); + } + } + + TllmGenFmhaKernel(TllmGenFmhaKernel const&) = delete; + TllmGenFmhaKernel& operator=(TllmGenFmhaKernel const&) = delete; + void loadKernels() { // Build a lookup map for all kernels. @@ -322,7 +441,7 @@ class TllmGenFmhaKernel // The number of CtasQ and CtasKv per sequence, Ctas in the Y dimension, and Ctas in the Z // dimension. - computeNumCtas(options, params.mMultiProcessorCount); + computeNumCtasSerialized(options, params.mMultiProcessorCount); if (shouldUseNvrtc(options)) { @@ -416,7 +535,7 @@ class TllmGenFmhaKernel updateFmhaOptions(options, optionsFromArgs); checkTrtllmSkipCorrThreshold(options, params.mSkipCorrThreshold); - auto [numCtasX, numCtasY, numCtasZ] = computeNumCtas(options, params.mMultiProcessorCount); + auto [numCtasX, numCtasY, numCtasZ] = computeNumCtasSerialized(options, params.mMultiProcessorCount); tg::CudaRunner::Grid grid{numCtasX, numCtasY, numCtasZ}; if (shouldUseNvrtc(options)) @@ -427,11 +546,14 @@ class TllmGenFmhaKernel populateJsonConfig(options, sstream); fmhaConfig.mGenCfgJsonStr = sstream.str(); + // Serialize with the foreground run() path: the export library caches + // compiled kernels in an unsynchronized global LRU cache. + std::lock_guard interfaceLock(getFmhaInterfaceMutex()); fmhaConfig.mExecPath = getExecPath().c_str(); fmhaConfig.mCtaDim = ctaDim; fmhaConfig.mGrid = grid; auto const compileStart = std::chrono::steady_clock::now(); - mFmhaInterface.generateAndCompileKernel(fmhaConfig); + recordJITCacheResult(mFmhaInterface.generateAndCompileKernel(fmhaConfig)); auto const compileElapsed = std::chrono::steady_clock::now() - compileStart; auto const compileElapsedMs = std::chrono::duration(compileElapsed).count(); if (compileElapsedMs > 1000.0) // FIXME: Change to return cache status from FmhaInterface @@ -442,6 +564,153 @@ class TllmGenFmhaKernel } } +public: + // The set of kernels a warmup sweep compiles only depends on this kernel + // instance (dtypes/SM/sage), the device the sweep runs on (modules load + // into that device's CUDA context, so an identical configuration on + // another device still needs its own sweep), and the scalar runner params + // hashed below, not on the calling layer. Every attention layer with the + // same configuration requests an identical sweep during the warmup + // forward, and (with eagle3 one-model speculative decoding) the target + // layers and the draft layers request sweeps that differ only in the + // spec-decoding fields. Every runner param that selectKernel() reads must + // be hashed here: sweeps with equal fingerprints are deduplicated, in sync + // mode as well as in async mode. + uint64_t jitWarmupFingerprint(RunnerParams const& params, int deviceId) const + { + auto combine = [](uint64_t seed, uint64_t value) + { return seed ^ (value + 0x9e3779b97f4a7c15ull + (seed << 6) + (seed >> 2)); }; + uint32_t skipSoftmaxBits = 0; + static_assert( + sizeof(skipSoftmaxBits) == sizeof(params.mSkipSoftmaxThresholdScaleFactor), "expect 32-bit float"); + std::memcpy(&skipSoftmaxBits, ¶ms.mSkipSoftmaxThresholdScaleFactor, sizeof(skipSoftmaxBits)); + uint32_t scaleQBits = 0; + std::memcpy(&scaleQBits, ¶ms.mScaleQ, sizeof(scaleQBits)); + uint64_t fingerprint = 0; + for (uint64_t value : + std::initializer_list{static_cast(deviceId), static_cast(mSM), + static_cast(mDtypeQ), static_cast(mDtypeK), static_cast(mDtypeV), + static_cast(mDtypeOut), static_cast(mNumEltsPerSageAttnBlkQ), + static_cast(mNumEltsPerSageAttnBlkK), static_cast(mNumEltsPerSageAttnBlkP), + static_cast(mNumEltsPerSageAttnBlkV), static_cast(params.mQkvLayout), + static_cast(params.mMaskType), static_cast(params.mKernelType), + static_cast(params.mTileScheduler), static_cast(params.mMultiCtasKvMode), + static_cast(params.mUseBlockSparseAttention), static_cast(params.mHeadDimQk), + static_cast(params.mHeadDimV), static_cast(params.mNumHeadsQ), + static_cast(params.mNumHeadsKv), static_cast(params.mNumHeadsQPerKv), + static_cast(params.mMaxSeqLenQ), static_cast(params.mUseGenKernelForPrefill), + static_cast(params.mJITWarmupMaxNumRequests), + static_cast(params.mJITWarmupMaxSeqLenQ), static_cast(params.mJITWarmupMaxSeqLenKv), + static_cast(params.mAttentionWindowSize), static_cast(params.mChunkedAttentionSize), + static_cast(params.mMaxNumPagesPerSeqKv), static_cast(params.mNumTokensPerPage), + static_cast(params.mNumPagesInMemPool), static_cast(params.mMultiProcessorCount), + static_cast(skipSoftmaxBits), static_cast(scaleQBits), + static_cast(params.mSparseAttention), static_cast(params.mSparseTopK), + static_cast(params.mIsSpecDecTree), + static_cast(params.mSpecDecodingTargetMaxGenLen), + static_cast(params.mDsv4EpilogueFusion.enabled)}) + { + fingerprint = combine(fingerprint, value); + } + return fingerprint; + } + +private: + using JITWarmupRegistry = TllmGenFmhaJitWarmupRegistry; + + // Process-wide; intentionally leaked (heap-allocated, never destroyed) so a + // background sweep can touch it at any point during process teardown. + static JITWarmupRegistry& getJITWarmupRegistry() + { + static auto* const sRegistry = new JITWarmupRegistry(); + return *sRegistry; + } + + // Tally of the export library's kernel-cache results (one per + // generateAndCompileKernel() call). Intentionally leaked. + static TllmGenFmhaJitCompileStats& getJITCompileStats() + { + static auto* const sStats = new TllmGenFmhaJitCompileStats(); + return *sStats; + } + + static void recordJITCacheResult(FmhaInterface::KernelCacheStatus status) + { + using CacheResult = TllmGenFmhaJitCompileStats::CacheResult; + CacheResult result = CacheResult::kUnknown; + if (status == FmhaInterface::KernelCacheStatus::CacheHit) + { + result = CacheResult::kHit; + } + else if (status == FmhaInterface::KernelCacheStatus::CacheMiss) + { + result = CacheResult::kMiss; + } + getJITCompileStats().record(result); + } + + // Re-runs a finished background sweep on the calling thread. A cache-hit loop + // (milliseconds) when the background pass compiled exactly the keys the + // foreground derivation produces; anything it missed, or a sweep whose + // background task failed, is compiled here and shows up as "JIT Warmup: + // Warmup for ... took ..." log lines inside the verify window. + void verifyJITWarmupSweep(JITWarmupRegistry::PendingSweep const& sweep) + { + TLLM_LOG_INFO("TRTLLM-Gen FMHA JIT warmup: foreground verify pass started (fingerprint=%llx).", + static_cast(sweep.mFingerprint)); + runJITWarmupGrid(sweep.mParams); + TLLM_LOG_INFO("TRTLLM-Gen FMHA JIT warmup: foreground verify pass finished (fingerprint=%llx).", + static_cast(sweep.mFingerprint)); + } + + // Foreground verify pass (async mode only): the first eager (non-capturing, + // non-warmup) NVRTC-path run() re-runs the sweeps this kernel object + // scheduled on the current device. drainAndVerifyAllJITWarmups() normally + // gets there first; this is the fallback for callers that never invoke it. + void maybeRunJITWarmupVerifyPass(RunnerParams const& liveParams) + { + if (getJITWarmupRegistry().numUnverified() == 0 || liveParams.mJITWarmup) + { + return; + } + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + TLLM_CUDA_CHECK(cudaStreamIsCapturing(liveParams.stream, &captureStatus)); + if (captureStatus != cudaStreamCaptureStatusNone) + { + return; + } + // Only sweeps this kernel object scheduled on the current device: the + // sweep re-runs through this object's kernel table and loads modules into + // the current device's context. + int const currentDevice = tensorrt_llm::common::getDevice(); + while (auto sweep = getJITWarmupRegistry().claimUnverified(currentDevice, this)) + { + verifyJITWarmupSweep(*sweep); + } + } + + // If CUDA graphs are being captured on this stream while the background JIT + // warmup thread still has work queued, wait for it to finish. The worker + // thread runs in the relaxed capture mode so its module loads cannot + // invalidate an ongoing capture, but draining here additionally guarantees + // that every launch captured into a graph resolves against a fully + // populated kernel cache. + void maybeWaitForAsyncJITWarmup(cudaStream_t stream) const + { + if (!asyncJITWarmupEnabled() || !asyncJITWarmupWorkerCreated().load(std::memory_order_acquire) + || getAsyncJITWarmupWorker().numPendingTasks() == 0) + { + return; + } + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + TLLM_CUDA_CHECK(cudaStreamIsCapturing(stream, &captureStatus)); + if (captureStatus != cudaStreamCaptureStatusNone) + { + TLLM_LOG_INFO("TRTLLM-Gen FMHA: draining background JIT warmup before capturing CUDA graphs."); + getAsyncJITWarmupWorker().drain(); + } + } + void runJITWarmupGridIfRequested(RunnerParams const& runnerParams) { if (!runnerParams.mJITWarmup || runnerParams.mKernelType != FmhaKernelType::Generation) @@ -465,6 +734,41 @@ class TllmGenFmhaKernel TLLM_CHECK_WITH_INFO(maxBatchSize > 0 && maxSeqLenKv > 0 && (!useGenKernelForPrefill || maxSeqLenQ > 0), "TRTLLM-Gen Fmha Warmup Param is invalid."); + // Snapshot the runner params: the warmup sweep only reads scalar fields + // and never dereferences device pointers or the stream. + auto warmupParams = runnerParams; + warmupParams.stream = nullptr; + int const deviceId = tensorrt_llm::common::getDevice(); + uint64_t const fingerprint = jitWarmupFingerprint(warmupParams, deviceId); + // Identical sweeps are deduplicated by scheduleJITWarmupSweep(): every + // layer of the warmup forward requests the same grid, and only the first + // request per distinct configuration compiles anything (the rest would + // be pure cache-hit loops). + auto const result = scheduleJITWarmupSweep(getJITWarmupRegistry(), + asyncJITWarmupEnabled() ? &getAsyncJITWarmupWorker() : nullptr, fingerprint, warmupParams, deviceId, this, + [this, fingerprint](RunnerParams const& params, int device) + { + TLLM_CUDA_CHECK(cudaSetDevice(device)); + runJITWarmupGrid(params); + TLLM_LOG_INFO("TRTLLM-Gen FMHA JIT warmup: generation-kernel grid sweep finished (fingerprint=%llx).", + static_cast(fingerprint)); + }); + if (result == JITWarmupScheduleResult::kQueued) + { + TLLM_LOG_INFO( + "TRTLLM-Gen FMHA JIT warmup: scheduled background compilation of the generation-kernel grid " + "(device=%d, maxBatchSize=%d, maxSeqLenQ=%d, maxSeqLenKv=%d, fingerprint=%llx).", + deviceId, maxBatchSize, maxSeqLenQ, maxSeqLenKv, static_cast(fingerprint)); + } + } + + void runJITWarmupGrid(RunnerParams const& runnerParams) + { + bool const useGenKernelForPrefill = runnerParams.mUseGenKernelForPrefill; + int const maxBatchSize = runnerParams.mJITWarmupMaxNumRequests; + int const maxSeqLenQ = runnerParams.mJITWarmupMaxSeqLenQ; + int const maxSeqLenKv = runnerParams.mJITWarmupMaxSeqLenKv; + std::vector batchSizeCandidates = makeWarmupCandidateSizes(kDefaultWarmupBatchSizeCandidates, maxBatchSize); // Use specified Q for generation, and use our Q grid for prefill @@ -503,6 +807,60 @@ class TllmGenFmhaKernel } public: + // Pre-capture / pre-readiness barrier (called from ModelEngine.warmup between + // the graph-shape warmup pass and graph capture): waits for every queued + // background sweep, then re-runs each finished-but-unverified sweep on the + // calling thread through the kernel object and device that scheduled it. + // Must be called with no CUDA graph capture in progress. Returns the number of + // sweeps verified here. A no-op unless async warmup is enabled. + static int64_t drainAndVerifyAllJITWarmups() + { + if (!asyncJITWarmupEnabled()) + { + return 0; + } + if (asyncJITWarmupWorkerCreated().load(std::memory_order_acquire)) + { + getAsyncJITWarmupWorker().drain(); + } + + struct RestoreDevice + { + int mDevice{tensorrt_llm::common::getDevice()}; + + ~RestoreDevice() + { + cudaSetDevice(mDevice); + } + } restoreDevice; + + int64_t numVerified = 0; + while (auto sweep = getJITWarmupRegistry().claimAnyUnverified()) + { + TLLM_CUDA_CHECK(cudaSetDevice(sweep->mDeviceId)); + sweep->mInstance->verifyJITWarmupSweep(*sweep); + ++numVerified; + } + return numVerified; + } + + // Number of kernel-cache misses, i.e. NVRTC compiles, reported by the export + // library in this process. A value that is stable across the first request + // after readiness proves no request-time JIT; the same key compiled again + // after an LRU eviction counts again. + static int64_t numJITCacheMisses() + { + return getJITCompileStats().numMisses(); + } + + // Number of generateAndCompileKernel() calls whose cache result the export + // library did not report. Non-zero means the miss count is not a complete + // measure of compilation. + static int64_t numJITUnknownCacheResults() + { + return getJITCompileStats().numUnknown(); + } + void run(RunnerParams const& params) { if (params.mMaxSeqLenQ == 0 || params.mBatchSize == 0 @@ -556,7 +914,7 @@ class TllmGenFmhaKernel // The number of CtasQ and CtasKv per sequence, Ctas in the Y dimension, and Ctas in the Z // dimension. - auto [numCtasX, numCtasY, numCtasZ] = computeNumCtas(options, params.mMultiProcessorCount); + auto [numCtasX, numCtasY, numCtasZ] = computeNumCtasSerialized(options, params.mMultiProcessorCount); // Set the launch grid size. tg::CudaRunner::Grid grid{numCtasX, numCtasY, numCtasZ}; @@ -580,6 +938,14 @@ class TllmGenFmhaKernel if (shouldUseNvrtc(options)) { + // If a CUDA graph capture is starting on this stream, make sure the + // background JIT warmup finished before capturing any FMHA launch. + maybeWaitForAsyncJITWarmup(params.stream); + // Re-run finished background sweeps once on the foreground thread + // (cache-hit loop when the background keys match; compiles any gap + // here, pre-readiness, instead of on the first real requests). + maybeRunJITWarmupVerifyPass(params); + // nvrtc path - uses mFmhaInterface member for kernel caching FmhaOptions nvrtcOptions = options; #ifdef TLLM_RUBIN_FEATURES @@ -596,11 +962,14 @@ class TllmGenFmhaKernel populateJsonConfig(nvrtcOptions, sstream); fmhaConfig.mGenCfgJsonStr = sstream.str(); + // Serialize with the async JIT warmup thread: the export library caches + // compiled kernels in an unsynchronized global LRU cache. + std::lock_guard interfaceLock(getFmhaInterfaceMutex()); fmhaConfig.mExecPath = getExecPath().c_str(); fmhaConfig.mCtaDim = ctaDim; fmhaConfig.mGrid = grid; auto const compileStart = std::chrono::steady_clock::now(); - mFmhaInterface.generateAndCompileKernel(fmhaConfig); + recordJITCacheResult(mFmhaInterface.generateAndCompileKernel(fmhaConfig)); auto const compileElapsed = std::chrono::steady_clock::now() - compileStart; auto const compileElapsedMs = std::chrono::duration(compileElapsed).count(); if (compileElapsedMs > 1000.0) // FIXME: Change to return cache status from FmhaInterface @@ -610,6 +979,24 @@ class TllmGenFmhaKernel "Possible JIT Cache Missing: TRTLLM-Gen FMHA generateAndCompileKernel took %.3f ms, kernelName=%s, " "batchSize=%d, maxSeqLenQ=%d, maxSeqLenKv=%d. This could affect performance measurement.", compileElapsedMs, kernelName.c_str(), params.mBatchSize, params.mMaxSeqLenQ, params.mMaxSeqLenKv); + // Dump the kernel-selection options so a cache miss can be diffed + // against what the warmup sweep compiled for the same kernel name. + TLLM_LOG_WARNING( + "JIT cache miss options: maskType=%d, kernelType=%d, tileScheduler=%d, multiCtasKvMode=%d, " + "headDimPerCtaV=%d, tileSizeQ=%d, tileSizeKv=%d, numTokensPerPage=%d, numHeadsQ=%d, " + "numHeadsKv=%d, numHeadsQPerKv=%d, numSpecDecodingTokens=%d, specDecodingTargetMaxGenLen=%d, " + "isCausalSpecDecodingGen=%d, isCustomSpecDecodingGen=%d, attentionWindowSize=%d, " + "chunkedAttentionSize=%d, sparseType=%d, skipsSoftmax=%d, clusterDimX=%d, reuseSmemKForV=%d, " + "numInstsQ=%d, numInstsKv=%d", + static_cast(options.mMaskType), static_cast(options.mFmhaKernelType), + static_cast(options.mTileScheduler), static_cast(options.mMultiCtasKvMode), + options.mHeadDimPerCtaV, options.mTileSizeQ, options.mTileSizeKv, options.mNumTokensPerPage, + options.mNumHeadsQ, options.mNumHeadsKv, options.mNumHeadsQPerKv, options.mNumSpecDecodingTokens, + options.mSpecDecodingTargetMaxGenLen, static_cast(options.mIsCausalSpecDecodingGen), + static_cast(options.mIsCustomSpecDecodingGen), options.mAttentionWindowSize, + options.mChunkedAttentionSize, static_cast(options.mSparseType), + static_cast(options.mSkipsSoftmaxWhenPossible), options.mClusterDimX, + static_cast(options.mReuseSmemKForV), options.mNumInstsQ, options.mNumInstsKv); } mFmhaInterface.run(fmhaConfig, fmhaData, params.stream, params.mMultiProcessorCount, 0); } @@ -664,12 +1051,25 @@ class TllmGenFmhaKernel } private: + // computeNumCtas queries cluster occupancy through the same shared dummy kernel that + // kernel selection mutates, so it is serialized on the same lock. Never call it while + // selectKernelAutotuned() holds the lock (std::mutex is not recursive). + std::tuple computeNumCtasSerialized( + FmhaOptions& options, int32_t multiProcessorCount) const + { + std::lock_guard selectionLock(getFmhaSelectionMutex()); + return computeNumCtas(options, multiProcessorCount); + } + // Autotuner entry for checkIfKernelExist, warmupOneKernel, and run. SM107 may select // CgaSmemReduction configs where clusterDimX * mMaxNumCtasPerSeqKv must stay <= 16; // the pre/post clamps below enforce that. Other architectures keep the legacy path. std::tuple selectKernelAutotuned( FmhaOptions options, FmhaOptionsFromArgs optionsFromArgs, int32_t multiProcessorCount) const { + // Kernel selection mutates and then queries the autotuner's shared dummy kernel; the + // background JIT warmup thread runs this same path. See getFmhaSelectionMutex(). + std::lock_guard selectionLock(getFmhaSelectionMutex()); // CGA SMEM reduction cluster limits are Rubin SM107-specific; clamping on SM100/SM103 // changed autotuner choices without fixing a known failure mode on those GPUs. if (mSM != kSM_107) @@ -841,9 +1241,15 @@ class TllmGenFmhaKernel info); } - std::string const& getExecPath() const +public: + // Install path handed to the export library for NVRTC compilation. Depends on + // no instance state. The string is intentionally leaked: a background warmup + // sweep may read it at any point up to process exit, whichever path created + // the worker (see fmhaJitWarmup.h). + static std::string const& getExecPath() { - static std::string execPathStr; + static auto* const execPathPtr = new std::string(); + std::string& execPathStr = *execPathPtr; if (execPathStr.empty()) { // Get build directory relative path from CMake macro @@ -974,6 +1380,7 @@ class TllmGenFmhaKernel return execPathStr; } +private: // Prepare pointers for TMA descriptors. static std::tuple getDevicePtrs( TllmGenFmhaRunnerParams const& runnerParams, int32_t bitsPerElt) diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index ac9053a76e18..c8aba5b851f6 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/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 diff --git a/cpp/tensorrt_llm/thop/trtllmGenFmhaJitWarmupOp.cpp b/cpp/tensorrt_llm/thop/trtllmGenFmhaJitWarmupOp.cpp new file mode 100644 index 000000000000..0b8bb41f8d04 --- /dev/null +++ b/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 + +#include + +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); +} diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index 95b7bdc7d596..28e54925a2e2 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/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 diff --git a/cpp/tests/unit_tests/kernels/fmhaJitWarmupFingerprintTest.cpp b/cpp/tests/unit_tests/kernels/fmhaJitWarmupFingerprintTest.cpp new file mode 100644 index 000000000000..17135619bcf7 --- /dev/null +++ b/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 + +#include + +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 makeKernel(Data_type dtypeOut = Data_type::DATA_TYPE_BF16) +{ + return std::make_unique(/*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(0x1); + withPointers.qPtr = reinterpret_cast(0x2); + withPointers.kvPtr = reinterpret_cast(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(0x4); + EXPECT_EQ(kernel->jitWarmupFingerprint(fusedWithBuffers, 0), kernel->jitWarmupFingerprint(fused, 0)); +} + +} // namespace diff --git a/cpp/tests/unit_tests/kernels/fmhaJitWarmupProcessExitTest.cpp b/cpp/tests/unit_tests/kernels/fmhaJitWarmupProcessExitTest.cpp new file mode 100644 index 000000000000..146de3b6e232 --- /dev/null +++ b/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 + +#include +#include +#include +#include +#include + +namespace +{ + +using tensorrt_llm::kernels::asyncJITWarmupWorkerCreated; +using tensorrt_llm::kernels::getAsyncJITWarmupWorker; +using tensorrt_llm::kernels::TllmGenFmhaKernel; + +std::atomic gTaskFinished{false}; +std::atomic 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 diff --git a/cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp b/cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp new file mode 100644 index 000000000000..696c4302c72d --- /dev/null +++ b/cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp @@ -0,0 +1,474 @@ +/* + * 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. + */ + +// Host-only tests for the asynchronous TRTLLM-Gen FMHA JIT warmup building +// blocks (fmhaJitWarmup.h): the background worker's lifecycle, the sweep +// registry's state machine, and the scheduling path the kernel runs them +// through. No GPU is involved; the "sweeps" here are plain callbacks. + +#include "tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace +{ + +using tensorrt_llm::kernels::JITWarmupScheduleResult; +using tensorrt_llm::kernels::JITWarmupState; +using tensorrt_llm::kernels::scheduleJITWarmupSweep; +using tensorrt_llm::kernels::TllmGenFmhaAsyncWarmupWorker; +using tensorrt_llm::kernels::TllmGenFmhaJitCompileStats; +using tensorrt_llm::kernels::TllmGenFmhaJitWarmupRegistry; + +// Stand-ins for the kernel object and the sweep parameters the registry stores. +struct FakeKernel +{ + int mId; +}; + +struct FakeParams +{ + int mMaxSeqLenQ; + int mMaxSeqLenKv; +}; + +using Registry = TllmGenFmhaJitWarmupRegistry; + +// A gate the test opens to let a queued task finish, so "pending" states can be +// observed deterministically. +class Gate +{ +public: + void open() + { + { + std::lock_guard lock(mMutex); + mOpen = true; + } + mCv.notify_all(); + } + + void wait() + { + std::unique_lock lock(mMutex); + mCv.wait(lock, [this]() { return mOpen; }); + } + +private: + std::mutex mMutex; + std::condition_variable mCv; + bool mOpen{false}; +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Worker +//////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(TllmGenFmhaAsyncWarmupWorker, RunsTasksInOrderAndDrainWaitsForAll) +{ + TllmGenFmhaAsyncWarmupWorker worker; + std::mutex mutex; + std::vector order; + Gate gate; + + ASSERT_TRUE(worker.enqueue( + [&]() + { + gate.wait(); + std::lock_guard lock(mutex); + order.push_back(1); + })); + ASSERT_TRUE(worker.enqueue( + [&]() + { + std::lock_guard lock(mutex); + order.push_back(2); + })); + EXPECT_EQ(worker.numPendingTasks(), 2); + + gate.open(); + worker.drain(); + EXPECT_EQ(worker.numPendingTasks(), 0); + EXPECT_EQ(order, (std::vector{1, 2})); +} + +TEST(TllmGenFmhaAsyncWarmupWorker, ThreadInitRunsOnWorkerThreadBeforeTasks) +{ + std::thread::id initThread; + std::thread::id taskThread; + TllmGenFmhaAsyncWarmupWorker worker([&]() { initThread = std::this_thread::get_id(); }); + ASSERT_TRUE(worker.enqueue([&]() { taskThread = std::this_thread::get_id(); })); + worker.drain(); + EXPECT_EQ(initThread, taskThread); + EXPECT_NE(initThread, std::this_thread::get_id()); +} + +TEST(TllmGenFmhaAsyncWarmupWorker, FailingTaskDoesNotStrandTheQueue) +{ + TllmGenFmhaAsyncWarmupWorker worker; + std::atomic secondRan{false}; + ASSERT_TRUE(worker.enqueue([]() { throw std::runtime_error("compile failed"); })); + ASSERT_TRUE(worker.enqueue([&]() { secondRan = true; })); + worker.drain(); + EXPECT_EQ(worker.numPendingTasks(), 0); + EXPECT_TRUE(secondRan); +} + +TEST(TllmGenFmhaAsyncWarmupWorker, ShutdownWithPendingWorkFinishesItThenRejectsNewWork) +{ + TllmGenFmhaAsyncWarmupWorker worker; + std::atomic numRan{0}; + Gate gate; + ASSERT_TRUE(worker.enqueue( + [&]() + { + gate.wait(); + ++numRan; + })); + ASSERT_TRUE(worker.enqueue([&]() { ++numRan; })); + EXPECT_FALSE(worker.isShutDown()); + + // shutdown() must wait for the queued work rather than drop it. + std::thread shutdownThread([&]() { worker.shutdown(); }); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + EXPECT_EQ(numRan, 0); + gate.open(); + shutdownThread.join(); + + EXPECT_EQ(numRan, 2); + EXPECT_EQ(worker.numPendingTasks(), 0); + EXPECT_TRUE(worker.isShutDown()); + + // Later work is not queued; the caller runs it inline. + EXPECT_FALSE(worker.enqueue([&]() { ++numRan; })); + EXPECT_EQ(numRan, 2); + // drain() after shutdown() returns immediately. + worker.drain(); + // shutdown() is idempotent. + worker.shutdown(); +} + +TEST(TllmGenFmhaAsyncWarmupWorker, ShutdownWithoutEverStartingIsSafe) +{ + TllmGenFmhaAsyncWarmupWorker worker; + worker.shutdown(); + EXPECT_TRUE(worker.isShutDown()); + EXPECT_FALSE(worker.enqueue([]() {})); + worker.drain(); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Registry +//////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(TllmGenFmhaJitWarmupRegistry, DeduplicatesIdenticalSweeps) +{ + Registry registry; + FakeKernel kernel{1}; + FakeParams params{4, 1024}; + EXPECT_FALSE(registry.tryRegister(0xA, params, JITWarmupState::kScheduled, /*deviceId=*/0, &kernel)); + EXPECT_TRUE(registry.tryRegister(0xA, params, JITWarmupState::kScheduled, /*deviceId=*/0, &kernel)); + EXPECT_EQ(registry.size(), 1u); + EXPECT_EQ(registry.state(0xA), JITWarmupState::kScheduled); +} + +TEST(TllmGenFmhaJitWarmupRegistry, SynchronousSweepsAreVerifiedOnRegistration) +{ + Registry registry; + FakeKernel kernel{1}; + EXPECT_FALSE(registry.tryRegister(0xA, FakeParams{4, 1024}, JITWarmupState::kVerified, 0, &kernel)); + EXPECT_EQ(registry.numUnverified(), 0); + EXPECT_FALSE(registry.claimUnverified(0, &kernel).has_value()); + EXPECT_FALSE(registry.claimAnyUnverified().has_value()); +} + +TEST(TllmGenFmhaJitWarmupRegistry, BackgroundDoneIsClaimedExactlyOnceByItsOwner) +{ + Registry registry; + FakeKernel kernel{1}; + FakeParams params{4, 1024}; + ASSERT_FALSE(registry.tryRegister(0xA, params, JITWarmupState::kScheduled, 0, &kernel)); + EXPECT_EQ(registry.numUnverified(), 0); + EXPECT_FALSE(registry.claimUnverified(0, &kernel).has_value()); + + registry.markBackgroundDone(0xA); + EXPECT_EQ(registry.numUnverified(), 1); + EXPECT_EQ(registry.state(0xA), JITWarmupState::kBackgroundDone); + + auto sweep = registry.claimUnverified(0, &kernel); + ASSERT_TRUE(sweep.has_value()); + EXPECT_EQ(sweep->mFingerprint, 0xAu); + EXPECT_EQ(sweep->mParams.mMaxSeqLenQ, 4); + EXPECT_EQ(sweep->mParams.mMaxSeqLenKv, 1024); + EXPECT_EQ(sweep->mDeviceId, 0); + EXPECT_EQ(sweep->mInstance, &kernel); + EXPECT_EQ(registry.numUnverified(), 0); + EXPECT_EQ(registry.state(0xA), JITWarmupState::kVerified); + + EXPECT_FALSE(registry.claimUnverified(0, &kernel).has_value()); + // A second markBackgroundDone (e.g. from a retried task) must not reopen it. + registry.markBackgroundDone(0xA); + EXPECT_EQ(registry.numUnverified(), 0); + EXPECT_EQ(registry.state(0xA), JITWarmupState::kVerified); +} + +TEST(TllmGenFmhaJitWarmupRegistry, SameDeviceMultiConfigIsolation) +{ + // Two kernel objects on one device (e.g. fp8 and bf16 attention configs). + // Each may only verify its own sweep. + Registry registry; + FakeKernel kernelA{1}; + FakeKernel kernelB{2}; + ASSERT_FALSE(registry.tryRegister(0xA, FakeParams{4, 1024}, JITWarmupState::kScheduled, 0, &kernelA)); + ASSERT_FALSE(registry.tryRegister(0xB, FakeParams{1, 2048}, JITWarmupState::kScheduled, 0, &kernelB)); + registry.markBackgroundDone(0xA); + registry.markBackgroundDone(0xB); + EXPECT_EQ(registry.numUnverified(), 2); + + // A reaches the verify pass first and must not take B's sweep. + auto sweepA = registry.claimUnverified(0, &kernelA); + ASSERT_TRUE(sweepA.has_value()); + EXPECT_EQ(sweepA->mInstance, &kernelA); + EXPECT_EQ(sweepA->mFingerprint, 0xAu); + EXPECT_FALSE(registry.claimUnverified(0, &kernelA).has_value()); + EXPECT_EQ(registry.state(0xB), JITWarmupState::kBackgroundDone); + EXPECT_EQ(registry.numUnverified(), 1); + + auto sweepB = registry.claimUnverified(0, &kernelB); + ASSERT_TRUE(sweepB.has_value()); + EXPECT_EQ(sweepB->mInstance, &kernelB); + EXPECT_EQ(sweepB->mFingerprint, 0xBu); + EXPECT_EQ(registry.numUnverified(), 0); +} + +TEST(TllmGenFmhaJitWarmupRegistry, MultiDeviceSweepsAreDistinctAndClaimedPerDevice) +{ + // The same configuration on two devices is two sweeps (compiled modules are + // per device context), and each device only verifies its own. + Registry registry; + FakeKernel kernelDev0{1}; + FakeKernel kernelDev1{2}; + FakeParams params{4, 1024}; + ASSERT_FALSE(registry.tryRegister(0xA0, params, JITWarmupState::kScheduled, 0, &kernelDev0)); + ASSERT_FALSE(registry.tryRegister(0xA1, params, JITWarmupState::kScheduled, 1, &kernelDev1)); + EXPECT_EQ(registry.size(), 2u); + registry.markBackgroundDone(0xA0); + registry.markBackgroundDone(0xA1); + + EXPECT_FALSE(registry.claimUnverified(1, &kernelDev0).has_value()); + EXPECT_FALSE(registry.claimUnverified(0, &kernelDev1).has_value()); + + auto sweep0 = registry.claimUnverified(0, &kernelDev0); + ASSERT_TRUE(sweep0.has_value()); + EXPECT_EQ(sweep0->mDeviceId, 0); + auto sweep1 = registry.claimUnverified(1, &kernelDev1); + ASSERT_TRUE(sweep1.has_value()); + EXPECT_EQ(sweep1->mDeviceId, 1); + EXPECT_EQ(registry.numUnverified(), 0); +} + +TEST(TllmGenFmhaJitWarmupRegistry, ClaimAnyUnverifiedServesTheDrainAndVerifyBarrier) +{ + Registry registry; + FakeKernel kernelA{1}; + FakeKernel kernelB{2}; + ASSERT_FALSE(registry.tryRegister(0xA, FakeParams{4, 1024}, JITWarmupState::kScheduled, 0, &kernelA)); + ASSERT_FALSE(registry.tryRegister(0xB, FakeParams{1, 2048}, JITWarmupState::kScheduled, 1, &kernelB)); + registry.markBackgroundDone(0xA); + registry.markBackgroundDone(0xB); + + int numClaimed = 0; + while (auto sweep = registry.claimAnyUnverified()) + { + // Each claimed sweep carries the owner and device the barrier must dispatch to. + EXPECT_TRUE((sweep->mInstance == &kernelA && sweep->mDeviceId == 0) + || (sweep->mInstance == &kernelB && sweep->mDeviceId == 1)); + ++numClaimed; + } + EXPECT_EQ(numClaimed, 2); + EXPECT_EQ(registry.numUnverified(), 0); + EXPECT_EQ(registry.state(0xA), JITWarmupState::kVerified); + EXPECT_EQ(registry.state(0xB), JITWarmupState::kVerified); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Compile stats +//////////////////////////////////////////////////////////////////////////////////////////////////// + +TEST(TllmGenFmhaJitCompileStats, CountsEveryMissIncludingRepeatsOfTheSameKey) +{ + using CacheResult = TllmGenFmhaJitCompileStats::CacheResult; + TllmGenFmhaJitCompileStats stats; + EXPECT_EQ(stats.numMisses(), 0); + EXPECT_EQ(stats.numHits(), 0); + EXPECT_EQ(stats.numUnknown(), 0); + + stats.record(CacheResult::kMiss); // first compile of a key + stats.record(CacheResult::kHit); // same key served from the cache + stats.record(CacheResult::kMiss); // same key compiled again after an LRU eviction + EXPECT_EQ(stats.numMisses(), 2); + EXPECT_EQ(stats.numHits(), 1); + EXPECT_EQ(stats.numUnknown(), 0); + + // An unreported result is neither a hit nor a miss; it must stay visible so a + // test cannot read "no misses" as "no compiles". + stats.record(CacheResult::kUnknown); + EXPECT_EQ(stats.numMisses(), 2); + EXPECT_EQ(stats.numUnknown(), 1); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Worker + registry together: the scheduling protocol used by fmhaKernels.h. +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Drives the kernel's scheduling path, scheduleJITWarmupSweep(), with the +// real worker and registry; only the sweep body is a stand-in. Sweeps are +// tallied by the thread they ran on, so a test can tell background runs from +// inline runs and verify-pass re-runs. +struct Scheduler +{ + // The registry is declared first so it outlives the worker's shutdown, which + // runs any still-queued task against it. + Registry mRegistry; + TllmGenFmhaAsyncWarmupWorker mWorker; + std::thread::id const mMainThread{std::this_thread::get_id()}; + std::atomic mBackgroundSweeps{0}; + std::atomic mForegroundSweeps{0}; + + JITWarmupScheduleResult schedule(uint64_t fingerprint, FakeParams const& params, FakeKernel* kernel, + bool failInBackground, bool async = true, int deviceId = 0) + { + return scheduleJITWarmupSweep(mRegistry, async ? &mWorker : nullptr, fingerprint, params, deviceId, kernel, + [this, failInBackground](FakeParams const&, int) + { + if (std::this_thread::get_id() == mMainThread) + { + ++mForegroundSweeps; + return; + } + if (failInBackground) + { + throw std::runtime_error("nvrtc failed"); + } + ++mBackgroundSweeps; + }); + } + + // The pre-capture barrier: drain, then re-run every finished sweep once. + int drainAndVerifyAll() + { + mWorker.drain(); + int numVerified = 0; + while (auto sweep = mRegistry.claimAnyUnverified()) + { + ++mForegroundSweeps; + ++numVerified; + } + return numVerified; + } +}; + +TEST(TllmGenFmhaJitWarmupProtocol, EveryScheduledSweepIsVerifiedExactlyOnceBeforeCapture) +{ + Scheduler scheduler; + FakeKernel kernel{1}; + EXPECT_EQ(scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, /*failInBackground=*/false), + JITWarmupScheduleResult::kQueued); + EXPECT_EQ(scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, false), JITWarmupScheduleResult::kDuplicate); + EXPECT_EQ(scheduler.schedule(0xB, FakeParams{1, 1024}, &kernel, false), JITWarmupScheduleResult::kQueued); + + EXPECT_EQ(scheduler.drainAndVerifyAll(), 2); + EXPECT_EQ(scheduler.mBackgroundSweeps, 2); + EXPECT_EQ(scheduler.mForegroundSweeps, 2); + EXPECT_EQ(scheduler.mRegistry.numUnverified(), 0); + // Nothing left for a later verify pass. + EXPECT_EQ(scheduler.drainAndVerifyAll(), 0); +} + +TEST(TllmGenFmhaJitWarmupProtocol, BackgroundFailureIsRetriedOnTheForeground) +{ + Scheduler scheduler; + FakeKernel kernel{1}; + EXPECT_EQ(scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, /*failInBackground=*/true), + JITWarmupScheduleResult::kQueued); + scheduler.mWorker.drain(); + // The failed sweep is not stranded in kScheduled: it is handed to the verify pass. + EXPECT_EQ(scheduler.mRegistry.state(0xA), JITWarmupState::kBackgroundDone); + EXPECT_EQ(scheduler.mBackgroundSweeps, 0); + + EXPECT_EQ(scheduler.drainAndVerifyAll(), 1); + EXPECT_EQ(scheduler.mForegroundSweeps, 1); + EXPECT_EQ(scheduler.mRegistry.state(0xA), JITWarmupState::kVerified); +} + +TEST(TllmGenFmhaJitWarmupProtocol, ShutDownWorkerFallsBackToInlineCompilation) +{ + Scheduler scheduler; + FakeKernel kernel{1}; + scheduler.mWorker.shutdown(); + EXPECT_EQ(scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, false), JITWarmupScheduleResult::kRanInline); + // Compiled inline by the caller, and still verified exactly once. + EXPECT_EQ(scheduler.mForegroundSweeps, 1); + EXPECT_EQ(scheduler.drainAndVerifyAll(), 1); + EXPECT_EQ(scheduler.mForegroundSweeps, 2); + EXPECT_EQ(scheduler.mBackgroundSweeps, 0); +} + +TEST(TllmGenFmhaJitWarmupProtocol, SyncModeRunsInlineAndStillDeduplicates) +{ + // Async warmup disabled: the sweep runs on the caller, is registered as + // verified, and an identical later request is dropped just as in async mode. + Scheduler scheduler; + FakeKernel kernel{1}; + EXPECT_EQ(scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, false, /*async=*/false), + JITWarmupScheduleResult::kRanInline); + EXPECT_EQ(scheduler.mForegroundSweeps, 1); + EXPECT_EQ(scheduler.mRegistry.state(0xA), JITWarmupState::kVerified); + EXPECT_EQ(scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, false, /*async=*/false), + JITWarmupScheduleResult::kDuplicate); + EXPECT_EQ(scheduler.mForegroundSweeps, 1); + // Nothing for the barrier to verify. + EXPECT_EQ(scheduler.drainAndVerifyAll(), 0); + EXPECT_EQ(scheduler.mBackgroundSweeps, 0); +} + +TEST(TllmGenFmhaJitWarmupProtocol, BackgroundSweepRunsOnItsOwnDevice) +{ + // The sweep body receives the device the request was registered on, not + // whatever is current on the worker thread. + Registry registry; + TllmGenFmhaAsyncWarmupWorker worker; + FakeKernel kernel{1}; + std::atomic seenDevice{-1}; + EXPECT_EQ(scheduleJITWarmupSweep(registry, &worker, 0xA, FakeParams{4, 1024}, /*deviceId=*/3, &kernel, + [&seenDevice](FakeParams const&, int deviceId) { seenDevice = deviceId; }), + JITWarmupScheduleResult::kQueued); + worker.drain(); + EXPECT_EQ(seenDevice.load(), 3); + auto const sweep = registry.claimUnverified(3, &kernel); + ASSERT_TRUE(sweep.has_value()); + EXPECT_EQ(sweep->mDeviceId, 3); +} + +} // namespace diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index a502b6ada99a..000592192a83 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -458,6 +458,11 @@ def _warmup_scheduled( # NVSHMEM). gc.collect() torch.cuda.empty_cache() + # Every background TRTLLM-Gen FMHA sweep was queued by the attention + # warmup above; finish and verify them before the first graph-shape + # forward, so no compile (or wait on one) lands inside a forward that + # runs MoE all-to-alls. + self._drain_and_verify_trtllm_gen_fmha_jit_warmup() with self._warmup_timer.phase("cuda_graph_capture"): self._graphs.warmup_and_capture(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") @@ -514,6 +519,38 @@ def _warmup_scheduled( return self._freeze_eager_workspace_floor() + def _drain_and_verify_trtllm_gen_fmha_jit_warmup(self) -> None: + """Wait for the background TRTLLM-Gen FMHA JIT warmup and verify it. + + Blocks until every background compilation sweep has finished, then + re-runs each finished sweep once on this thread: a cache-hit loop when + the background pass compiled every kernel, and the place where any + kernel it missed (or a sweep whose background task failed) is compiled + instead of on the first real request. Runs whether or not CUDA graphs + are enabled, so configurations without graphs still finish + verification before readiness. A no-op unless async warmup is enabled + (``TRTLLM_GEN_FMHA_ASYNC_WARMUP=1``). + """ + if not issubclass(self._config.attention_backend.Metadata, TrtllmAttentionMetadata): + return + if not torch.ops.trtllm.trtllm_gen_fmha_async_jit_warmup_enabled(): + return + with self._warmup_timer.phase("attention_jit_drain_and_verify"): + num_verified = torch.ops.trtllm.trtllm_gen_fmha_jit_warmup_drain_and_verify() + if num_verified: + logger.info( + f"TRTLLM-Gen FMHA JIT warmup: verified {num_verified} " + "background warmup sweep(s) before CUDA graph capture." + ) + # Each rank only waited for its own compilation. The graph-shape + # warmup and capture forwards that follow run MoE all-to-alls, the + # capture pass under the serving timeouts, so a rank that enters + # one while a peer is still compiling could spend that timeout + # waiting. Meet the peers here first. + allgather = self._agreement.allgather() + if allgather is not None: + allgather(1) + def _freeze_eager_workspace_floor(self) -> EagerWorkspaceReclaimer | None: if not self._eager_workspace_reclaimable: return None diff --git a/tests/unittest/_torch/attention/test_trtllm_gen_fmha_async_jit_warmup.py b/tests/unittest/_torch/attention/test_trtllm_gen_fmha_async_jit_warmup.py new file mode 100644 index 000000000000..42538ec55129 --- /dev/null +++ b/tests/unittest/_torch/attention/test_trtllm_gen_fmha_async_jit_warmup.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end check of the asynchronous TRTLLM-Gen FMHA JIT warmup. + +With ``TRTLLM_GEN_FMHA_ASYNC_WARMUP=1`` the warmup grid compiles on a +background thread, and the engine's pre-capture barrier waits for it and +re-verifies every sweep before CUDA graphs are captured. This test boots a +small model with CUDA graphs enabled and checks, from the engine process: + +* capture succeeds (the LLM reaches readiness), +* the warmup exercised the NVRTC path at all (the count of distinct kernel + configurations requested from the FMHA export library grew during warmup), +* the first requests after readiness compile nothing: that count is unchanged + across them. The export library compiles a configuration the first time it + is requested, so this is a direct measure of compilation activity, not a + latency heuristic. + +The flag is read once per process, so each arm runs in a fresh subprocess +with the worker in that process (``TLLM_WORKER_USE_SINGLE_PROCESS=1``), where +the counters are observable. The flag-off arm is the control: the same +invariant must hold for the synchronous warmup. +""" + +import json +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest +from utils.llm_data import llm_models_root +from utils.util import skip_pre_blackwell + +_CHILD = textwrap.dedent(""" + import json + import sys + + import torch + + from tensorrt_llm import LLM, SamplingParams + from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig + + model_dir = sys.argv[1] + misses = torch.ops.trtllm.trtllm_gen_fmha_jit_num_cache_misses + unknown = torch.ops.trtllm.trtllm_gen_fmha_jit_num_unknown_cache_results + result = {"misses_before": misses()} + llm = LLM( + model=model_dir, + attn_backend="TRTLLM", + tensor_parallel_size=1, + max_batch_size=8, + max_num_tokens=512, + cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2, 4, 8]), + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.3), + ) + with llm: + result["async_enabled"] = bool( + torch.ops.trtllm.trtllm_gen_fmha_async_jit_warmup_enabled()) + result["misses_ready"] = misses() + prompts = [ + "The capital of France is", + "Write one sentence about the ocean.", + "List three prime numbers:", + "Explain why the sky is blue in a few words.", + ] + outputs = llm.generate(prompts, SamplingParams(max_tokens=48)) + result["num_outputs"] = len(outputs) + result["all_generated"] = all( + len(o.outputs[0].token_ids) > 0 for o in outputs) + result["misses_after_first_requests"] = misses() + result["unknown_cache_results"] = unknown() + print("RESULT " + json.dumps(result)) +""") + + +def _run_arm(model_dir: Path, async_warmup: bool) -> dict: + env = dict(os.environ) + env["TRTLLM_GEN_FMHA_ASYNC_WARMUP"] = "1" if async_warmup else "0" + # Run the executor worker in the child's own process so the kernel-key + # counter reflects the engine that served the requests. + env["TLLM_WORKER_USE_SINGLE_PROCESS"] = "1" + proc = subprocess.run( + [sys.executable, "-c", _CHILD, str(model_dir)], + env=env, + capture_output=True, + text=True, + timeout=1200, + ) + assert proc.returncode == 0, ( + f"engine subprocess failed (async_warmup={async_warmup}):\n" + f"--- stdout ---\n{proc.stdout[-4000:]}\n" + f"--- stderr ---\n{proc.stderr[-4000:]}" + ) + lines = [line for line in proc.stdout.splitlines() if line.startswith("RESULT ")] + assert lines, f"no RESULT line in child stdout:\n{proc.stdout[-4000:]}" + return json.loads(lines[-1][len("RESULT ") :]) + + +@pytest.fixture(scope="module") +def model_dir() -> Path: + return llm_models_root() / "Qwen3/Qwen3-0.6B" + + +@skip_pre_blackwell +@pytest.mark.parametrize("async_warmup", [True, False], ids=["async_warmup", "sync_warmup"]) +def test_first_requests_after_readiness_do_not_jit_compile( + model_dir: Path, async_warmup: bool +) -> None: + result = _run_arm(model_dir, async_warmup) + + assert result["async_enabled"] is async_warmup + assert result["num_outputs"] == 4 and result["all_generated"] + + # Every compile request must have come back with a reported cache result; + # otherwise the miss count is not a complete measure and the invariants + # below could hold while compiles went uncounted. + assert result["unknown_cache_results"] == 0, ( + f"{result['unknown_cache_results']} compile request(s) returned an " + "unknown cache status; the miss count cannot be trusted" + ) + + # The warmup must have gone through the NVRTC path, otherwise the + # invariant below would hold vacuously. + assert result["misses_ready"] > result["misses_before"], ( + "warmup reported no kernel-cache misses; this model/config does not " + "exercise the TRTLLM-Gen JIT path" + ) + + # No cache miss after readiness, i.e. the first requests compiled nothing: + # every kernel they needed was compiled (and, for the async arm, verified) + # before capture. A key evicted and recompiled would count as a miss too. + assert result["misses_after_first_requests"] == result["misses_ready"], ( + f"{result['misses_after_first_requests'] - result['misses_ready']} " + "kernel(s) were compiled by the first requests after readiness" + ) diff --git a/tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py b/tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py new file mode 100644 index 000000000000..46d7af67b963 --- /dev/null +++ b/tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py @@ -0,0 +1,229 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ordering tests for the TRTLLM-Gen FMHA JIT warmup barrier in decoder warmup. + +The barrier (``torch.ops.trtllm.trtllm_gen_fmha_jit_warmup_drain_and_verify`` +plus the rank sync) must run after every sweep has been queued and before +the graph-shape warmup and capture passes, so that no compile (or wait on a +background compile) lands inside a forward that runs MoE all-to-alls, and so +that every attention kernel the warmup derived is compiled before readiness. +These tests drive ``DecoderWarmup.run`` with the graph passes, the MoE +serving-timeout switch and the ops replaced by recorders; the C++ state +machine behind the ops is covered by +``cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp``. +""" + +import contextlib +from collections.abc import Iterator +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import graphs as graphs_module +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder import warmup as warmup_module +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.graphs import DecoderCudaGraphs +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.warmup import ( + DecoderWarmup, + WarmupAgreement, +) +from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer + +pytestmark = pytest.mark.cpu_only + + +class _Recorder: + """Records the order of barrier calls, rank syncs, timeout switches and graph passes.""" + + def __init__(self, drain_result: int = 0, async_enabled: bool = True): + self.events: list[str] = [] + self.drain_result = drain_result + self.async_enabled = async_enabled + + def drain_and_verify(self) -> int: + self.events.append("drain_and_verify") + return self.drain_result + + def allgather(self, value): + self.events.append("rank_sync") + return [value, value] + + @contextlib.contextmanager + def serving_timeouts(self) -> Iterator[None]: + self.events.append("serving_timeouts_enter") + try: + yield + finally: + self.events.append("serving_timeouts_exit") + + +class _GraphRunnerStub: + def __init__(self) -> None: + self.is_warmup_only = False + self.padding_dummy_requests: dict[str, object] = {"stale": object()} + + @contextlib.contextmanager + def allow_capture(self) -> Iterator[None]: + yield + + def preallocate_padding_dummies(self, resource_manager: object) -> None: + pass + + +class _RecordingGraphs(DecoderCudaGraphs): + """The real two-pass ``warmup_and_capture`` over a recording graph pass.""" + + def __init__(self, recorder: _Recorder) -> None: + self.cuda_graph_runner = _GraphRunnerStub() + self._recorder = recorder + + def run_cuda_graph_warmup(self, resource_manager: object) -> None: + self._recorder.events.append( + "graph_warmup" if self.cuda_graph_runner.is_warmup_only else "graph_capture" + ) + + +def _engine( + recorder: _Recorder, *, trtllm_attention: bool = True, world_size: int = 1 +) -> DecoderWarmup: + """Build a warmup engine carrying only what the scheduled path reads. + + The context warmups are off, which skips the attention, general, autotuner + and Mamba phases, so the graph-shape warmup and capture passes are the only + forwards. ``world_size > 1`` gives the forward group a recording allgather. + """ + dist = ( + SimpleNamespace(world_size=world_size, allgather=recorder.allgather) + if world_size > 1 + else None + ) + mapping = SimpleNamespace(cp_size=1, tp_size=1, dwdp_enabled=False, has_cp_helix=lambda: False) + engine = object.__new__(DecoderWarmup) + engine._agreement = WarmupAgreement(dist, mapping) + engine.dist = dist + engine.mapping = mapping + engine._config = SimpleNamespace( + attention_backend=SimpleNamespace( + Metadata=TrtllmAttentionMetadata if trtllm_attention else AttentionMetadata + ) + ) + engine._warmup_timer = _WarmupTimer(rank=0) + engine._metrics = {} + engine._graphs = _RecordingGraphs(recorder) + engine._inputs = SimpleNamespace(guided_decoder=None) + engine._metadata = SimpleNamespace(attn_metadata=None) + engine._context_warmups_supported = False + engine._eager_workspace_reclaimable = False + engine.kv_cache_manager_key = "kv_cache_manager" + # Model-specific prewarms the scheduled path runs unconditionally; each + # needs a model or a device to decide it has nothing to do. + engine._prewarm_engram_kernels = mock.Mock() + engine._prewarm_csa2_sparse_indexer = mock.Mock() + engine._prewarm_cute_dsl_indexer_q = mock.Mock() + engine._ensure_dsa_attn_metadata_for_warmup = mock.Mock() + return engine + + +def _resource_manager() -> SimpleNamespace: + kv_cache_manager = SimpleNamespace() + return SimpleNamespace(get_resource_manager=lambda key: kv_cache_manager) + + +@contextlib.contextmanager +def _patched(recorder: _Recorder) -> Iterator[None]: + with ( + mock.patch.object(warmup_module, "warmup_sampling_module"), + mock.patch.object(warmup_module, "log_mem_snapshot"), + mock.patch.object(warmup_module, "AutoTuner"), + mock.patch.object( + graphs_module, "moe_comm_serving_timeouts", side_effect=recorder.serving_timeouts + ), + mock.patch.object( + torch.ops.trtllm, + "trtllm_gen_fmha_jit_warmup_drain_and_verify", + create=True, + side_effect=recorder.drain_and_verify, + ), + mock.patch.object( + torch.ops.trtllm, + "trtllm_gen_fmha_async_jit_warmup_enabled", + create=True, + side_effect=lambda: recorder.async_enabled, + ), + ): + yield + + +# The graph-shape warmup pass keeps the warmup MoE timeouts; only the capture +# pass switches to the serving ones. +_GRAPH_PASSES = ["graph_warmup", "serving_timeouts_enter", "graph_capture", "serving_timeouts_exit"] + + +def test_barrier_runs_before_the_graph_shape_warmup_and_capture() -> None: + recorder = _Recorder() + engine = _engine(recorder) + with _patched(recorder): + engine.run(_resource_manager()) + + assert recorder.events == ["drain_and_verify", *_GRAPH_PASSES] + assert engine._graphs.cuda_graph_runner.is_warmup_only is False + assert engine._graphs.cuda_graph_runner.padding_dummy_requests == {} + + +def test_barrier_phase_is_not_nested_in_the_capture_phase() -> None: + # Both phases are recorded, so the warmup summary must not count the + # drain twice (once on its own and once inside cuda_graph_capture). + recorder = _Recorder() + engine = _engine(recorder) + with _patched(recorder): + engine.run(_resource_manager()) + timings = engine._warmup_timer.timings + assert {"attention_jit_drain_and_verify", "cuda_graph_capture"} <= set(timings) + assert recorder.events.index("drain_and_verify") < recorder.events.index("graph_warmup") + + +def test_barrier_reports_foreground_recovery_of_background_failures() -> None: + # A background sweep that failed is re-run by the barrier; the op reports + # how many sweeps it had to verify on the foreground, and warmup proceeds + # to the graph passes only after that. + recorder = _Recorder(drain_result=2) + engine = _engine(recorder) + with _patched(recorder), mock.patch.object(warmup_module, "logger") as warmup_logger: + engine.run(_resource_manager()) + assert recorder.events == ["drain_and_verify", *_GRAPH_PASSES] + logged = " ".join(str(arg) for call in warmup_logger.info.call_args_list for arg in call.args) + assert "verified 2 background warmup sweep(s)" in logged + + +def test_barrier_is_skipped_for_non_trtllm_attention_backends() -> None: + recorder = _Recorder() + engine = _engine(recorder, trtllm_attention=False) + with _patched(recorder): + engine.run(_resource_manager()) + assert recorder.events == _GRAPH_PASSES + + +def test_barrier_is_skipped_when_async_warmup_is_disabled() -> None: + # Flag off (the default): the synchronous warmup already compiled + # everything inline, so there is nothing to drain, verify or sync on. + recorder = _Recorder(async_enabled=False) + engine = _engine(recorder) + with _patched(recorder): + engine.run(_resource_manager()) + assert recorder.events == _GRAPH_PASSES + + +def test_ranks_meet_after_the_local_drain_and_before_the_graph_passes() -> None: + # Each rank only waits for its own compilation; a rank entering the graph + # passes while a peer is still compiling would wait on that peer inside an + # all-to-all, so the ranks meet after draining and before the first pass. + recorder = _Recorder() + engine = _engine(recorder, world_size=2) + with _patched(recorder): + engine.run(_resource_manager()) + # The leading sync is the phase-entry agreement that precedes every + # warmup phase; the one after the drain is the barrier's. + assert recorder.events == ["rank_sync", "drain_and_verify", "rank_sync", *_GRAPH_PASSES]