From 62b9b93683ac3226b6c828e20036a54e5b101c36 Mon Sep 17 00:00:00 2001 From: Aaron Meoded Date: Tue, 1 Sep 2026 11:46:16 -0400 Subject: [PATCH 01/11] [None][perf] Overlap TRTLLM-Gen FMHA JIT warmup compilation with engine warmup (opt-in) At engine startup, TRTLLM-Gen FMHA generation kernels are NVRTC-compiled synchronously inside the first warmup forward (~5-11s per kernel variant, 80-160s total for MLA fp8-KV configs, ~2min for sparse gen kernels; scales with node CPU speed and now with the widened seqLenQ warmup candidates). With TRTLLM_GEN_FMHA_ASYNC_WARMUP=1 the warmup grid compiles on a single background thread, overlapping autotuning, CUDA-graph capture, and memory profiling. Default off; flag-off behavior is unchanged apart from deduplicating identical per-layer warmup sweeps. Safety: - All calls into the FMHA export lib are serialized behind a global interface mutex: the lib's kernelCache is an unsynchronized global LRU (hits mutate it too), and kernel selection mutates shared state via the FmhaAutoTuner occupancy probe (attribute set->query on a shared dummy kernel). - The background queue is drained before CUDA-graph capture; run() drains if invoked on a capturing stream. - A foreground verify pass re-walks each sweep before readiness, so any gap degrades to inline compilation - never a post-ready miss. - Worker-thread start failure and background compile exceptions degrade to inline compilation. Validated on B200 deployments (DeepSeek-family MLA fp8-KV, eagle3 one-model spec decode): critical-path warmup 163s -> ~26-30s, spec-decode acceptance and steady-state throughput unchanged, no first-traffic compile misses. Signed-off-by: Aaron Meoded --- .../trtllmGenKernels/fmha/fmhaKernels.h | 533 ++++++++++++++++-- 1 file changed, 492 insertions(+), 41 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 63fe47c0b604..61096d1212e5 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -19,18 +19,26 @@ #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 "tensorrt_llm/common/cudaDriverWrapper.h" @@ -95,6 +103,171 @@ 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. +inline std::mutex& getFmhaInterfaceMutex() +{ + static std::mutex sMutex; + 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. +inline std::mutex& getFmhaSelectionMutex() +{ + static std::mutex sMutex; + 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; +} + +// A single background worker thread that runs the JIT warmup grid compilation +// off the startup critical path. A single thread is intentional: all compile +// calls serialize on getFmhaInterfaceMutex() anyway (the export library's +// kernel cache is not thread-safe), so a pool would add no parallelism. +class TllmGenFmhaAsyncWarmupWorker +{ +public: + // Enqueue for the background worker; returns false if the worker thread could + // not be started, in which case the task was NOT queued and the caller must + // run it inline. + bool enqueue(std::function task) + { + { + std::lock_guard lock(mMutex); + mQueue.push_back(std::move(task)); + mNumPendingTasks.fetch_add(1, std::memory_order_acq_rel); + if (!mThreadStarted) + { + // Only mark started after the thread object is successfully constructed: + // if std::thread throws (resource exhaustion), the queued task must not + // be stranded behind a worker that never runs (drain() would deadlock). + try + { + std::thread( + [this]() + { + // Relaxed capture mode: CUDA calls made by this 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."); + } + workerLoop(); + }) + .detach(); + mThreadStarted = true; + } + catch (std::exception const& e) + { + mQueue.pop_back(); + mNumPendingTasks.fetch_sub(1, std::memory_order_acq_rel); + TLLM_LOG_WARNING( + "TRTLLM-Gen FMHA async JIT warmup: worker thread failed to start (%s); compiling inline.", + e.what()); + return false; + } + } + } + mQueueCv.notify_one(); + return true; + } + + int numPendingTasks() const + { + return mNumPendingTasks.load(std::memory_order_acquire); + } + + // Blocks until all enqueued warmup tasks have finished. + void drain() + { + std::unique_lock lock(mMutex); + mDrainCv.wait(lock, [this]() { return mNumPendingTasks.load(std::memory_order_acquire) == 0; }); + } + +private: + void workerLoop() + { + while (true) + { + std::function task; + { + std::unique_lock lock(mMutex); + mQueueCv.wait(lock, [this]() { return !mQueue.empty(); }); + 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::mutex mMutex; + std::condition_variable mQueueCv; + std::condition_variable mDrainCv; + std::deque> mQueue; + std::atomic mNumPendingTasks{0}; + bool mThreadStarted{false}; +}; + +inline TllmGenFmhaAsyncWarmupWorker& getAsyncJITWarmupWorker() +{ + // Intentionally leaked: the detached worker thread parks on the queue + // condition variable forever and must outlive static destruction. + static TllmGenFmhaAsyncWarmupWorker* sWorker = new TllmGenFmhaAsyncWarmupWorker(); + return *sWorker; +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + class TllmGenFmhaKernel { @@ -260,17 +433,20 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); - std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); - - // Check if the options are valid or not. - checkFmhaOptions(options, optionsFromArgs); - // Update the options if needed. - updateFmhaOptions(options, optionsFromArgs); - - // The number of CtasQ and CtasKv per sequence, Ctas in the Y dimension, and Ctas in the Z - // dimension. - computeNumCtas(options, params.mMultiProcessorCount); + { + std::lock_guard selectionLock(getFmhaSelectionMutex()); + FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); + + // Check if the options are valid or not. + checkFmhaOptions(options, optionsFromArgs); + // Update the options if needed. + updateFmhaOptions(options, optionsFromArgs); + + // The number of CtasQ and CtasKv per sequence, Ctas in the Y dimension, and Ctas in the Z + // dimension. + computeNumCtas(options, params.mMultiProcessorCount); + } if (shouldUseNvrtc(options)) { @@ -357,13 +533,17 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); - std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); + int numCtasX = 0, numCtasY = 0, numCtasZ = 0; + { + std::lock_guard selectionLock(getFmhaSelectionMutex()); + FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); - checkFmhaOptions(options, optionsFromArgs); - updateFmhaOptions(options, optionsFromArgs); + checkFmhaOptions(options, optionsFromArgs); + updateFmhaOptions(options, optionsFromArgs); - auto [numCtasX, numCtasY, numCtasZ] = computeNumCtas(options, params.mMultiProcessorCount); + std::tie(numCtasX, numCtasY, numCtasZ) = computeNumCtas(options, params.mMultiProcessorCount); + } tg::CudaRunner::Grid grid{numCtasX, numCtasY, numCtasZ}; if (shouldUseNvrtc(options)) @@ -374,6 +554,9 @@ 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; @@ -389,6 +572,180 @@ class TllmGenFmhaKernel } } + // The set of kernels a warmup sweep compiles only depends on this kernel + // instance (dtypes/SM/sage) 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. + uint64_t jitWarmupFingerprint(RunnerParams const& params) 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(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)}) + { + fingerprint = combine(fingerprint, value); + } + return fingerprint; + } + + // Per-configuration warmup progress. + enum class JITWarmupState : int + { + kScheduled = 0, // Background sweep enqueued (async mode). + kBackgroundDone, // Background sweep finished; foreground verify pass pending. + kVerified, // Sweep ran (or re-ran) on a foreground thread. + }; + + struct JITWarmupConfig + { + RunnerParams mParams; + JITWarmupState mState; + }; + + static std::mutex& getJITWarmupRegistryMutex() + { + static std::mutex sMutex; + return sMutex; + } + + static std::unordered_map& getJITWarmupRegistry() + { + static auto* sRegistry = new std::unordered_map(); + return *sRegistry; + } + + // Number of configs in kBackgroundDone state (cheap gate for the verify pass). + static std::atomic& getNumUnverifiedJITWarmups() + { + static std::atomic sCount{0}; + return sCount; + } + + // Registers a warmup configuration; returns true if an identical sweep was + // already executed or scheduled. + bool tryRegisterJITWarmupConfig(RunnerParams const& params, JITWarmupState initialState) const + { + uint64_t const fingerprint = jitWarmupFingerprint(params); + std::lock_guard lock(getJITWarmupRegistryMutex()); + auto& registry = getJITWarmupRegistry(); + if (registry.find(fingerprint) != registry.end()) + { + return true; + } + JITWarmupConfig config; + config.mParams = params; + config.mState = initialState; + registry.emplace(fingerprint, std::move(config)); + return false; + } + + static void markJITWarmupBackgroundDone(uint64_t fingerprint) + { + std::lock_guard lock(getJITWarmupRegistryMutex()); + auto it = getJITWarmupRegistry().find(fingerprint); + if (it != getJITWarmupRegistry().end() && it->second.mState == JITWarmupState::kScheduled) + { + it->second.mState = JITWarmupState::kBackgroundDone; + getNumUnverifiedJITWarmups().fetch_add(1, std::memory_order_acq_rel); + } + } + + // Foreground verify pass (async mode only). After a background sweep + // finishes, the first eager (non-capturing, non-warmup) NVRTC-path run() + // re-executes the same sweep on a foreground thread. If the background pass + // compiled exactly the keys the foreground derivation produces, this is a + // pure cache-hit loop (milliseconds). Any divergence — e.g. from a source of + // nondeterminism this patch has not closed — is compiled here, before + // readiness, instead of stalling the first real requests, and shows up as + // "JIT Warmup: Warmup for ... took ..." log lines inside the verify window. + void maybeRunJITWarmupVerifyPass(RunnerParams const& liveParams) + { + if (getNumUnverifiedJITWarmups().load(std::memory_order_acquire) == 0 || liveParams.mJITWarmup) + { + return; + } + cudaStreamCaptureStatus captureStatus = cudaStreamCaptureStatusNone; + TLLM_CUDA_CHECK(cudaStreamIsCapturing(liveParams.stream, &captureStatus)); + if (captureStatus != cudaStreamCaptureStatusNone) + { + return; + } + while (true) + { + uint64_t fingerprint = 0; + RunnerParams sweepParams{}; + { + std::lock_guard lock(getJITWarmupRegistryMutex()); + auto& registry = getJITWarmupRegistry(); + auto it = std::find_if(registry.begin(), registry.end(), [](auto const& entry) + { return entry.second.mState == JITWarmupState::kBackgroundDone; }); + if (it == registry.end()) + { + return; + } + fingerprint = it->first; + sweepParams = it->second.mParams; + it->second.mState = JITWarmupState::kVerified; + getNumUnverifiedJITWarmups().fetch_sub(1, std::memory_order_acq_rel); + } + TLLM_LOG_INFO("TRTLLM-Gen FMHA JIT warmup: foreground verify pass started (fingerprint=%llx).", + static_cast(fingerprint)); + runJITWarmupGrid(sweepParams); + TLLM_LOG_INFO("TRTLLM-Gen FMHA JIT warmup: foreground verify pass finished (fingerprint=%llx).", + static_cast(fingerprint)); + } + } + + // 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() || 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) @@ -412,6 +769,67 @@ class TllmGenFmhaKernel TLLM_CHECK_WITH_INFO(maxBatchSize > 0 && maxSeqLenKv > 0 && (!useGenKernelForPrefill || maxSeqLenQ > 0), "TRTLLM-Gen Fmha Warmup Param is invalid."); + // Deduplicate identical sweeps: 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). + // 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; + if (tryRegisterJITWarmupConfig( + warmupParams, asyncJITWarmupEnabled() ? JITWarmupState::kScheduled : JITWarmupState::kVerified)) + { + return; + } + + if (asyncJITWarmupEnabled()) + { + uint64_t const fingerprint = jitWarmupFingerprint(warmupParams); + int const deviceId = tensorrt_llm::common::getDevice(); + TLLM_LOG_INFO( + "TRTLLM-Gen FMHA JIT warmup: scheduling background compilation of the generation-kernel grid " + "(device=%d, maxBatchSize=%d, maxSeqLenQ=%d, maxSeqLenKv=%d, fingerprint=%llx).", + deviceId, maxBatchSize, maxSeqLenQ, maxSeqLenKv, static_cast(fingerprint)); + bool const queued = getAsyncJITWarmupWorker().enqueue( + [this, warmupParams, deviceId, fingerprint]() + { + try + { + TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); + runJITWarmupGrid(warmupParams); + } + catch (...) + { + // Mark done even on failure so the fingerprint is not stranded in + // kScheduled: the foreground verify pass then re-runs this sweep + // and compiles whatever the background pass did not finish, + // instead of dedupe silently suppressing it forever. + markJITWarmupBackgroundDone(fingerprint); + throw; + } + markJITWarmupBackgroundDone(fingerprint); + TLLM_LOG_INFO("TRTLLM-Gen FMHA JIT warmup: background compilation of the generation-kernel grid " + "finished (fingerprint=%llx).", + static_cast(fingerprint)); + }); + if (!queued) + { + runJITWarmupGrid(runnerParams); + markJITWarmupBackgroundDone(fingerprint); + } + return; + } + + runJITWarmupGrid(runnerParams); + } + + 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 @@ -465,35 +883,39 @@ class TllmGenFmhaKernel parseOptionsFromRunnerParams(params, options); options.mCudaArch = intToCudaArch(mSM); - FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); - std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); - - // Overwrite AutoTuner decision: SageAttention with SfsPV is known to cause regression to persistent scheduler. - // Remove this overwritten once we refresh the cubin kernels that containing the related fix. - if (mNumEltsPerSageAttnBlkP + mNumEltsPerSageAttnBlkV > 0) + int numCtasX = 0, numCtasY = 0, numCtasZ = 0; { - options.mTileScheduler = TileScheduler::Static; - } + std::lock_guard selectionLock(getFmhaSelectionMutex()); + FmhaAutoTuner autoTuner(options, optionsFromArgs, params.mMultiProcessorCount); + std::tie(options, optionsFromArgs, ctaDim) = autoTuner.selectKernel(); - // Check if the options are valid or not. - checkFmhaOptions(options, optionsFromArgs); - // Update the options if needed. - updateFmhaOptions(options, optionsFromArgs); + // Overwrite AutoTuner decision: SageAttention with SfsPV is known to cause regression to persistent + // scheduler. Remove this overwritten once we refresh the cubin kernels that containing the related fix. + if (mNumEltsPerSageAttnBlkP + mNumEltsPerSageAttnBlkV > 0) + { + options.mTileScheduler = TileScheduler::Static; + } - // Any caller that selects MultiCtasKvMode must supply the partial-reduction scratch pool - // and per-CTA counter; fail fast here instead of silently falling back to Disabled. - if (options.mMultiCtasKvMode == tensorrt_llm::kernels::MultiCtasKvMode::GmemReduction - || options.mMultiCtasKvMode == tensorrt_llm::kernels::MultiCtasKvMode::GmemReductionWithSeparateKernel) - { - TLLM_CHECK_WITH_INFO(params.multiCtasKvScratchPtr != nullptr && params.multiCtasKvCounterPtr != nullptr, - "MultiCtasKvScratchPtr/MultiCtasKvCounterPtr must be non-null when fmha kernel uses gmem-based " - "multi-CTA reduction. " - "The dispatcher must allocate and pass these buffers."); - } + // Check if the options are valid or not. + checkFmhaOptions(options, optionsFromArgs); + // Update the options if needed. + updateFmhaOptions(options, optionsFromArgs); - // 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); + // Any caller that selects MultiCtasKvMode must supply the partial-reduction scratch pool + // and per-CTA counter; fail fast here instead of silently falling back to Disabled. + if (options.mMultiCtasKvMode == tensorrt_llm::kernels::MultiCtasKvMode::GmemReduction + || options.mMultiCtasKvMode == tensorrt_llm::kernels::MultiCtasKvMode::GmemReductionWithSeparateKernel) + { + TLLM_CHECK_WITH_INFO(params.multiCtasKvScratchPtr != nullptr && params.multiCtasKvCounterPtr != nullptr, + "MultiCtasKvScratchPtr/MultiCtasKvCounterPtr must be non-null when fmha kernel uses gmem-based " + "multi-CTA reduction. " + "The dispatcher must allocate and pass these buffers."); + } + + // The number of CtasQ and CtasKv per sequence, Ctas in the Y dimension, and Ctas in the Z + // dimension. + std::tie(numCtasX, numCtasY, numCtasZ) = computeNumCtas(options, params.mMultiProcessorCount); + } // Set the launch grid size. tg::CudaRunner::Grid grid{numCtasX, numCtasY, numCtasZ}; @@ -516,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 FmhaConfig fmhaConfig; fmhaConfig.mOptions = options; @@ -523,6 +953,9 @@ class TllmGenFmhaKernel populateJsonConfig(options, 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; @@ -537,6 +970,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); } From 59821ba08bfb5c9f798152adb9e672cbdfb39069 Mon Sep 17 00:00:00 2001 From: Aaron Meoded Date: Tue, 1 Sep 2026 13:30:15 -0400 Subject: [PATCH 02/11] [None][fix] Address review: per-device warmup registry, drain on kernel destruction - Include the current device in the JIT warmup fingerprint and store it in the registry entry; the verify pass now only consumes sweeps registered for the current device. Compiled modules load into the calling device's CUDA context, so an identical configuration on another device still needs its own sweep (single-process multi-GPU flows). - Drain the async warmup worker in ~TllmGenFmhaKernel(): a queued background task captures 'this', and nothing forced a drain before static teardown of the owning factory if a process exits mid-startup. The worker singleton is intentionally leaked, so draining during static destruction is safe. Also delete copy operations (the factory owns the kernel via unique_ptr; nothing copies it). Signed-off-by: Aaron Meoded --- .../trtllmGenKernels/fmha/fmhaKernels.h | 42 ++++++++++++++++--- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 61096d1212e5..e6d51fa1d971 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -321,6 +321,23 @@ class TllmGenFmhaKernel { } + ~TllmGenFmhaKernel() + { + // A queued background warmup task captures `this`. Scheduled sweeps drain + // before CUDA-graph capture in normal operation, but nothing forces that + // before this object is destroyed (e.g. static teardown of the owning + // factory while startup is still in flight). The worker singleton is + // intentionally leaked, so draining here is safe even during static + // destruction. + if (asyncJITWarmupEnabled()) + { + getAsyncJITWarmupWorker().drain(); + } + } + + TllmGenFmhaKernel(TllmGenFmhaKernel const&) = delete; + TllmGenFmhaKernel& operator=(TllmGenFmhaKernel const&) = delete; + void loadKernels() { // Build a lookup map for all kernels. @@ -573,8 +590,9 @@ class TllmGenFmhaKernel } // The set of kernels a warmup sweep compiles only depends on this kernel - // instance (dtypes/SM/sage) and the scalar runner params hashed below, not - // on the calling layer. Every attention layer with the same configuration + // instance (dtypes/SM/sage), the current device (modules load into the + // calling device's CUDA context), 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. @@ -589,7 +607,12 @@ class TllmGenFmhaKernel uint32_t scaleQBits = 0; std::memcpy(&scaleQBits, ¶ms.mScaleQ, sizeof(scaleQBits)); uint64_t fingerprint = 0; - for (uint64_t value : std::initializer_list{static_cast(mSM), + // Include the current device: compiled modules are loaded into the calling + // device's CUDA context, so an identical configuration on another device + // still needs its own sweep. Warmup requests always run with the target + // device current. + for (uint64_t value : std::initializer_list{ + static_cast(tensorrt_llm::common::getDevice()), static_cast(mSM), static_cast(mDtypeQ), static_cast(mDtypeK), static_cast(mDtypeV), static_cast(mDtypeOut), static_cast(mNumEltsPerSageAttnBlkQ), static_cast(mNumEltsPerSageAttnBlkK), static_cast(mNumEltsPerSageAttnBlkP), @@ -629,6 +652,9 @@ class TllmGenFmhaKernel { RunnerParams mParams; JITWarmupState mState; + // Device the sweep was registered for. The fingerprint already includes it; + // kept here so the verify pass only picks up sweeps for the current device. + int mDeviceId; }; static std::mutex& getJITWarmupRegistryMutex() @@ -664,6 +690,7 @@ class TllmGenFmhaKernel JITWarmupConfig config; config.mParams = params; config.mState = initialState; + config.mDeviceId = tensorrt_llm::common::getDevice(); registry.emplace(fingerprint, std::move(config)); return false; } @@ -699,6 +726,9 @@ class TllmGenFmhaKernel { return; } + // Only verify sweeps registered for the current device: the sweep re-runs on + // this thread, and compiled modules load into the current device's context. + int const currentDevice = tensorrt_llm::common::getDevice(); while (true) { uint64_t fingerprint = 0; @@ -706,8 +736,10 @@ class TllmGenFmhaKernel { std::lock_guard lock(getJITWarmupRegistryMutex()); auto& registry = getJITWarmupRegistry(); - auto it = std::find_if(registry.begin(), registry.end(), [](auto const& entry) - { return entry.second.mState == JITWarmupState::kBackgroundDone; }); + auto it = std::find_if(registry.begin(), registry.end(), [currentDevice](auto const& entry) { + return entry.second.mState == JITWarmupState::kBackgroundDone + && entry.second.mDeviceId == currentDevice; + }); if (it == registry.end()) { return; From c32cc978e5dc834dda71884fdc15744823d7438b Mon Sep 17 00:00:00 2001 From: Aaron Meoded Date: Wed, 2 Sep 2026 12:48:41 -0400 Subject: [PATCH 03/11] [None][chore] Apply clang-format Signed-off-by: Aaron Meoded --- .../trtllmGenKernels/fmha/fmhaKernels.h | 89 ++++++++++--------- 1 file changed, 45 insertions(+), 44 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index e6d51fa1d971..b3b6216b31b3 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -172,22 +172,22 @@ class TllmGenFmhaAsyncWarmupWorker try { std::thread( - [this]() - { - // Relaxed capture mode: CUDA calls made by this 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) + [this]() { - 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."); - } - workerLoop(); - }) + // Relaxed capture mode: CUDA calls made by this 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."); + } + workerLoop(); + }) .detach(); mThreadStarted = true; } @@ -611,29 +611,28 @@ class TllmGenFmhaKernel // device's CUDA context, so an identical configuration on another device // still needs its own sweep. Warmup requests always run with the target // device current. - for (uint64_t value : std::initializer_list{ - static_cast(tensorrt_llm::common::getDevice()), 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)}) + for (uint64_t value : + std::initializer_list{static_cast(tensorrt_llm::common::getDevice()), + 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)}) { fingerprint = combine(fingerprint, value); } @@ -736,10 +735,11 @@ class TllmGenFmhaKernel { std::lock_guard lock(getJITWarmupRegistryMutex()); auto& registry = getJITWarmupRegistry(); - auto it = std::find_if(registry.begin(), registry.end(), [currentDevice](auto const& entry) { - return entry.second.mState == JITWarmupState::kBackgroundDone - && entry.second.mDeviceId == currentDevice; - }); + auto it = std::find_if(registry.begin(), registry.end(), + [currentDevice](auto const& entry) { + return entry.second.mState == JITWarmupState::kBackgroundDone + && entry.second.mDeviceId == currentDevice; + }); if (it == registry.end()) { return; @@ -840,8 +840,9 @@ class TllmGenFmhaKernel throw; } markJITWarmupBackgroundDone(fingerprint); - TLLM_LOG_INFO("TRTLLM-Gen FMHA JIT warmup: background compilation of the generation-kernel grid " - "finished (fingerprint=%llx).", + TLLM_LOG_INFO( + "TRTLLM-Gen FMHA JIT warmup: background compilation of the generation-kernel grid " + "finished (fingerprint=%llx).", static_cast(fingerprint)); }); if (!queued) From b72ad826662ed78d4d783ef332c74c6e66881948 Mon Sep 17 00:00:00 2001 From: Aaron Meoded Date: Tue, 22 Sep 2026 16:04:49 +0800 Subject: [PATCH 04/11] [None][fix] Per-kernel-instance warmup registry isolation; leak registry singletons Review feedback (chienchunhung): 1. The warmup registry is shared across all TllmGenFmhaKernel objects on the same device (different dtypes / attention configs). The verify pass filtered only by device, so kernel A could consume kernel B's kBackgroundDone entry, mark it verified, and leave B to JIT-compile on its first real request. Now each registry entry stores the registering kernel instance (mKernelInstance = this) and the verify pass filters on both device and instance. 2. The registry mutex, registry map, and unverified counter were function-local statics, destroyed before the factory's kernels during static teardown. A destructor drain() could then wait on a worker task using already-destroyed state. All three are now leaked heap singletons (same pattern as the worker itself), guaranteeing they outlive every kernel object. Signed-off-by: Aaron Meoded --- .../trtllmGenKernels/fmha/fmhaKernels.h | 35 ++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index b3b6216b31b3..292bd2e0337c 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -654,12 +654,25 @@ class TllmGenFmhaKernel // Device the sweep was registered for. The fingerprint already includes it; // kept here so the verify pass only picks up sweeps for the current device. int mDeviceId; + // The kernel instance that registered this sweep. The registry is shared + // across all TllmGenFmhaKernel objects on the same device (different dtypes + // / attention configs); without this, one instance's verify pass could + // consume another's entry and mark it verified without ever checking it. + TllmGenFmhaKernel const* mKernelInstance; }; + // All three singletons below are intentionally leaked (heap-allocated, never + // destroyed). The async warmup worker is also leaked, and a queued background + // task can use any of these at any time during process teardown. C++ destroys + // function-local statics in reverse construction order, so a non-leaked + // static mutex or counter could be destroyed before the factory's kernels, + // leaving the destructor's drain() waiting on a task that touches destroyed + // state. Leaking keeps them alive for the remainder of the process. + static std::mutex& getJITWarmupRegistryMutex() { - static std::mutex sMutex; - return sMutex; + static auto* sMutex = new std::mutex(); + return *sMutex; } static std::unordered_map& getJITWarmupRegistry() @@ -671,8 +684,8 @@ class TllmGenFmhaKernel // Number of configs in kBackgroundDone state (cheap gate for the verify pass). static std::atomic& getNumUnverifiedJITWarmups() { - static std::atomic sCount{0}; - return sCount; + static auto* sCount = new std::atomic(0); + return *sCount; } // Registers a warmup configuration; returns true if an identical sweep was @@ -690,6 +703,7 @@ class TllmGenFmhaKernel config.mParams = params; config.mState = initialState; config.mDeviceId = tensorrt_llm::common::getDevice(); + config.mKernelInstance = this; registry.emplace(fingerprint, std::move(config)); return false; } @@ -725,9 +739,13 @@ class TllmGenFmhaKernel { return; } - // Only verify sweeps registered for the current device: the sweep re-runs on - // this thread, and compiled modules load into the current device's context. + // Only verify sweeps registered for the current device AND kernel + // instance: the sweep re-runs on this thread through this kernel's + // table, and compiled modules load into the current device's context. + // Without the instance filter, one kernel (e.g. fp8) could consume + // another's (e.g. bf16) entry and mark it verified without checking it. int const currentDevice = tensorrt_llm::common::getDevice(); + TllmGenFmhaKernel const* const currentInstance = this; while (true) { uint64_t fingerprint = 0; @@ -736,9 +754,10 @@ class TllmGenFmhaKernel std::lock_guard lock(getJITWarmupRegistryMutex()); auto& registry = getJITWarmupRegistry(); auto it = std::find_if(registry.begin(), registry.end(), - [currentDevice](auto const& entry) { + [currentDevice, currentInstance](auto const& entry) { return entry.second.mState == JITWarmupState::kBackgroundDone - && entry.second.mDeviceId == currentDevice; + && entry.second.mDeviceId == currentDevice + && entry.second.mKernelInstance == currentInstance; }); if (it == registry.end()) { From 68bcdcee905973c25d451d7b3829ca3d286ffa72 Mon Sep 17 00:00:00 2001 From: Aaron Meoded Date: Wed, 7 Oct 2026 15:08:39 -0400 Subject: [PATCH 05/11] [None][fix] Pre-capture drain-and-verify barrier, ordered worker shutdown, per-sweep verification Review feedback (chienchunhung): - Pre-capture barrier: ModelEngine.warmup now runs torch.ops.trtllm.trtllm_gen_fmha_jit_warmup_drain_and_verify() between the graph-shape warmup pass and graph capture, outside the graph-enabled and general-warmup conditions. The op waits for every queued background sweep and re-runs each finished-but-unverified sweep on the calling thread, through the kernel object and device that scheduled it, so every kernel the warmup derived is compiled before any launch is captured and before readiness. The in-run() verify pass stays as the fallback for callers that never reach the barrier. - Ordered shutdown: the worker thread is now joinable and shut down through a std::atexit handler registered when the worker is first created (from the first warmup forward). atexit handlers run before the destruction of any static whose initialization completed earlier, so an in-flight sweep finishes while everything it reads is alive. getExecPath()'s static string, which the sweep would otherwise construct lazily on the worker thread, is touched on the registering thread first. The interface and selection mutexes are leaked like the registry. - The worker and the sweep registry move to fmhaJitWarmup.h, with no CUDA dependency, so the state machine is covered by host unit tests. Each registry entry records the owning kernel object and device; a verify pass only claims its own sweeps. - torch.ops.trtllm.trtllm_gen_fmha_jit_num_requested_kernel_keys() exposes the number of distinct NVRTC kernel configurations requested from the export library (hashed as its cache keys them), so a test can check for request-time compilation directly instead of inferring it from latency. - Merged with main: kernel selection and computeNumCtas now take the selection lock inside selectKernelAutotuned()/computeNumCtasSerialized() rather than at each call site. Signed-off-by: Aaron Meoded --- .../trtllmGenKernels/fmha/fmhaJitWarmup.h | 341 +++++++++++++++++ .../trtllmGenKernels/fmha/fmhaKernels.h | 358 +++++++----------- cpp/tensorrt_llm/thop/CMakeLists.txt | 1 + .../thop/trtllmGenFmhaJitWarmupOp.cpp | 67 ++++ .../_torch/pyexecutor/model_engine.py | 50 ++- 5 files changed, 585 insertions(+), 232 deletions(-) create mode 100644 cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h create mode 100644 cpp/tensorrt_llm/thop/trtllmGenFmhaJitWarmupOp.cpp 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..f87b76e1772a --- /dev/null +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h @@ -0,0 +1,341 @@ +/* + * 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 through a +// std::atexit handler registered when the worker is first created, i.e. from +// the first warmup forward. atexit handlers run before the destruction of any +// static object whose initialization completed earlier, so a sweep still in +// flight at process exit finishes while everything it reads (the export +// library's globals, the kernel objects, the mutexes, the install-path string) +// is still alive. The one caveat is a static that the sweep itself constructs +// lazily on the worker thread: such a static would be destroyed before the +// handler runs. fmhaKernels.h therefore touches those on the registering +// thread before the first task is queued. +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}; +}; + +} // 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 db197823fc39..929934d42004 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -32,6 +32,7 @@ #include #include #include +#include #include #include #include @@ -43,6 +44,7 @@ #include "tensorrt_llm/common/cudaDriverWrapper.h" #include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/kernels/multiHeadAttentionCommon.h" @@ -118,11 +120,12 @@ constexpr bool isSMCompatible(int gpuSM, int kernelSM) // 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. +// 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 std::mutex sMutex; - return sMutex; + static auto* const sMutex = new std::mutex(); + return *sMutex; } // Serializes kernel selection (FmhaAutoTuner::selectKernel and @@ -136,11 +139,12 @@ inline std::mutex& getFmhaInterfaceMutex() // 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. +// while getFmhaInterfaceMutex() is taken. Intentionally leaked, like the +// interface mutex. inline std::mutex& getFmhaSelectionMutex() { - static std::mutex sMutex; - return sMutex; + static auto* const sMutex = new std::mutex(); + return *sMutex; } // Whether the TRTLLM-Gen FMHA JIT warmup grid is compiled on a background @@ -156,121 +160,35 @@ inline bool asyncJITWarmupEnabled() return sEnabled; } -// A single background worker thread that runs the JIT warmup grid compilation -// off the startup critical path. A single thread is intentional: all compile -// calls serialize on getFmhaInterfaceMutex() anyway (the export library's -// kernel cache is not thread-safe), so a pool would add no parallelism. -class TllmGenFmhaAsyncWarmupWorker +// Process-wide worker for the background JIT warmup sweeps. See +// fmhaJitWarmup.h for the lifetime and shutdown-ordering contract. +inline TllmGenFmhaAsyncWarmupWorker& getAsyncJITWarmupWorker() { -public: - // Enqueue for the background worker; returns false if the worker thread could - // not be started, in which case the task was NOT queued and the caller must - // run it inline. - bool enqueue(std::function task) + static TllmGenFmhaAsyncWarmupWorker* const sWorker = []() { - { - std::lock_guard lock(mMutex); - mQueue.push_back(std::move(task)); - mNumPendingTasks.fetch_add(1, std::memory_order_acq_rel); - if (!mThreadStarted) + auto* worker = new TllmGenFmhaAsyncWarmupWorker( + []() { - // Only mark started after the thread object is successfully constructed: - // if std::thread throws (resource exhaustion), the queued task must not - // be stranded behind a worker that never runs (drain() would deadlock). - try - { - std::thread( - [this]() - { - // Relaxed capture mode: CUDA calls made by this 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."); - } - workerLoop(); - }) - .detach(); - mThreadStarted = true; - } - catch (std::exception const& e) + // 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) { - mQueue.pop_back(); - mNumPendingTasks.fetch_sub(1, std::memory_order_acq_rel); TLLM_LOG_WARNING( - "TRTLLM-Gen FMHA async JIT warmup: worker thread failed to start (%s); compiling inline.", - e.what()); - return false; + "TRTLLM-Gen FMHA async JIT warmup: failed to set relaxed stream-capture mode; " + "run() will drain this worker before any CUDA graph capture."); } - } - } - mQueueCv.notify_one(); - return true; - } - - int numPendingTasks() const - { - return mNumPendingTasks.load(std::memory_order_acquire); - } - - // Blocks until all enqueued warmup tasks have finished. - void drain() - { - std::unique_lock lock(mMutex); - mDrainCv.wait(lock, [this]() { return mNumPendingTasks.load(std::memory_order_acquire) == 0; }); - } - -private: - void workerLoop() - { - while (true) - { - std::function task; - { - std::unique_lock lock(mMutex); - mQueueCv.wait(lock, [this]() { return !mQueue.empty(); }); - 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::mutex mMutex; - std::condition_variable mQueueCv; - std::condition_variable mDrainCv; - std::deque> mQueue; - std::atomic mNumPendingTasks{0}; - bool mThreadStarted{false}; -}; - -inline TllmGenFmhaAsyncWarmupWorker& getAsyncJITWarmupWorker() -{ - // Intentionally leaked: the detached worker thread parks on the queue - // condition variable forever and must outlive static destruction. - static TllmGenFmhaAsyncWarmupWorker* sWorker = new TllmGenFmhaAsyncWarmupWorker(); + }); + // The worker is first created from the first warmup forward, after every + // static a sweep depends on has been constructed, so this handler drains + // and joins the worker before any of them is destroyed. The object itself + // is intentionally leaked. + std::atexit([]() { getAsyncJITWarmupWorker().shutdown(); }); + return worker; + }(); return *sWorker; } @@ -350,12 +268,11 @@ class TllmGenFmhaKernel ~TllmGenFmhaKernel() { - // A queued background warmup task captures `this`. Scheduled sweeps drain - // before CUDA-graph capture in normal operation, but nothing forces that - // before this object is destroyed (e.g. static teardown of the owning - // factory while startup is still in flight). The worker singleton is - // intentionally leaked, so draining here is safe even during static - // destruction. + // A queued background warmup task captures `this`. At process exit the + // worker's atexit handler has already drained and joined it (see + // fmhaJitWarmup.h); draining here as well covers a kernel object destroyed + // while the process keeps running. The worker is leaked, so this is safe + // during static destruction too. if (asyncJITWarmupEnabled()) { getAsyncJITWarmupWorker().drain(); @@ -625,6 +542,7 @@ class TllmGenFmhaKernel fmhaConfig.mGrid = grid; auto const compileStart = std::chrono::steady_clock::now(); mFmhaInterface.generateAndCompileKernel(fmhaConfig); + recordJITKernelKey(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 @@ -685,97 +603,66 @@ class TllmGenFmhaKernel return fingerprint; } - // Per-configuration warmup progress. - enum class JITWarmupState : int - { - kScheduled = 0, // Background sweep enqueued (async mode). - kBackgroundDone, // Background sweep finished; foreground verify pass pending. - kVerified, // Sweep ran (or re-ran) on a foreground thread. - }; + using JITWarmupRegistry = TllmGenFmhaJitWarmupRegistry; - struct JITWarmupConfig + // Process-wide; intentionally leaked (heap-allocated, never destroyed) so a + // background sweep can touch it at any point during process teardown. + static JITWarmupRegistry& getJITWarmupRegistry() { - RunnerParams mParams; - JITWarmupState mState; - // Device the sweep was registered for. The fingerprint already includes it; - // kept here so the verify pass only picks up sweeps for the current device. - int mDeviceId; - // The kernel instance that registered this sweep. The registry is shared - // across all TllmGenFmhaKernel objects on the same device (different dtypes - // / attention configs); without this, one instance's verify pass could - // consume another's entry and mark it verified without ever checking it. - TllmGenFmhaKernel const* mKernelInstance; - }; - - // All three singletons below are intentionally leaked (heap-allocated, never - // destroyed). The async warmup worker is also leaked, and a queued background - // task can use any of these at any time during process teardown. C++ destroys - // function-local statics in reverse construction order, so a non-leaked - // static mutex or counter could be destroyed before the factory's kernels, - // leaving the destructor's drain() waiting on a task that touches destroyed - // state. Leaking keeps them alive for the remainder of the process. + static auto* const sRegistry = new JITWarmupRegistry(); + return *sRegistry; + } - static std::mutex& getJITWarmupRegistryMutex() + // Keys of every NVRTC kernel configuration requested from the export library + // so far, hashed exactly as its kernel cache keys them: a key seen for the + // first time is a compile. Intentionally leaked. Guarded by + // getFmhaInterfaceMutex(), which both request sites already hold. + static std::unordered_set& getRequestedJITKernelKeys() { - static auto* sMutex = new std::mutex(); - return *sMutex; + static auto* const sKeys = new std::unordered_set(); + return *sKeys; } - static std::unordered_map& getJITWarmupRegistry() + // Must be called with getFmhaInterfaceMutex() held, after generateAndCompileKernel(). + static void recordJITKernelKey(FmhaConfig const& fmhaConfig) { - static auto* sRegistry = new std::unordered_map(); - return *sRegistry; + getRequestedJITKernelKeys().insert(std::hash{}(fmhaConfig.mOptions)); } - // Number of configs in kBackgroundDone state (cheap gate for the verify pass). - static std::atomic& getNumUnverifiedJITWarmups() + // Registers a warmup sweep for this kernel object on the current device; + // returns true if an identical sweep was already registered. + bool tryRegisterJITWarmupConfig(RunnerParams const& params, JITWarmupState initialState) { - static auto* sCount = new std::atomic(0); - return *sCount; + return getJITWarmupRegistry().tryRegister( + jitWarmupFingerprint(params), params, initialState, tensorrt_llm::common::getDevice(), this); } - // Registers a warmup configuration; returns true if an identical sweep was - // already executed or scheduled. - bool tryRegisterJITWarmupConfig(RunnerParams const& params, JITWarmupState initialState) const + static void markJITWarmupBackgroundDone(uint64_t fingerprint) { - uint64_t const fingerprint = jitWarmupFingerprint(params); - std::lock_guard lock(getJITWarmupRegistryMutex()); - auto& registry = getJITWarmupRegistry(); - if (registry.find(fingerprint) != registry.end()) - { - return true; - } - JITWarmupConfig config; - config.mParams = params; - config.mState = initialState; - config.mDeviceId = tensorrt_llm::common::getDevice(); - config.mKernelInstance = this; - registry.emplace(fingerprint, std::move(config)); - return false; + getJITWarmupRegistry().markBackgroundDone(fingerprint); } - static void markJITWarmupBackgroundDone(uint64_t fingerprint) + // 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) { - std::lock_guard lock(getJITWarmupRegistryMutex()); - auto it = getJITWarmupRegistry().find(fingerprint); - if (it != getJITWarmupRegistry().end() && it->second.mState == JITWarmupState::kScheduled) - { - it->second.mState = JITWarmupState::kBackgroundDone; - getNumUnverifiedJITWarmups().fetch_add(1, std::memory_order_acq_rel); - } + 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). After a background sweep - // finishes, the first eager (non-capturing, non-warmup) NVRTC-path run() - // re-executes the same sweep on a foreground thread. If the background pass - // compiled exactly the keys the foreground derivation produces, this is a - // pure cache-hit loop (milliseconds). Any divergence — e.g. from a source of - // nondeterminism this patch has not closed — is compiled here, before - // readiness, instead of stalling the first real requests, and shows up as - // "JIT Warmup: Warmup for ... took ..." log lines inside the verify window. + // 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 (getNumUnverifiedJITWarmups().load(std::memory_order_acquire) == 0 || liveParams.mJITWarmup) + if (getJITWarmupRegistry().numUnverified() == 0 || liveParams.mJITWarmup) { return; } @@ -785,40 +672,13 @@ class TllmGenFmhaKernel { return; } - // Only verify sweeps registered for the current device AND kernel - // instance: the sweep re-runs on this thread through this kernel's - // table, and compiled modules load into the current device's context. - // Without the instance filter, one kernel (e.g. fp8) could consume - // another's (e.g. bf16) entry and mark it verified without checking it. + // 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(); - TllmGenFmhaKernel const* const currentInstance = this; - while (true) + while (auto sweep = getJITWarmupRegistry().claimUnverified(currentDevice, this)) { - uint64_t fingerprint = 0; - RunnerParams sweepParams{}; - { - std::lock_guard lock(getJITWarmupRegistryMutex()); - auto& registry = getJITWarmupRegistry(); - auto it = std::find_if(registry.begin(), registry.end(), - [currentDevice, currentInstance](auto const& entry) { - return entry.second.mState == JITWarmupState::kBackgroundDone - && entry.second.mDeviceId == currentDevice - && entry.second.mKernelInstance == currentInstance; - }); - if (it == registry.end()) - { - return; - } - fingerprint = it->first; - sweepParams = it->second.mParams; - it->second.mState = JITWarmupState::kVerified; - getNumUnverifiedJITWarmups().fetch_sub(1, std::memory_order_acq_rel); - } - TLLM_LOG_INFO("TRTLLM-Gen FMHA JIT warmup: foreground verify pass started (fingerprint=%llx).", - static_cast(fingerprint)); - runJITWarmupGrid(sweepParams); - TLLM_LOG_INFO("TRTLLM-Gen FMHA JIT warmup: foreground verify pass finished (fingerprint=%llx).", - static_cast(fingerprint)); + verifyJITWarmupSweep(*sweep); } } @@ -881,6 +741,10 @@ class TllmGenFmhaKernel if (asyncJITWarmupEnabled()) { + // The install-path string is a function-local static that the sweep reads + // on the worker thread. Construct it on this thread, before the worker and + // its atexit shutdown exist (lifetime note in fmhaJitWarmup.h). + (void) getExecPath(); uint64_t const fingerprint = jitWarmupFingerprint(warmupParams); int const deviceId = tensorrt_llm::common::getDevice(); TLLM_LOG_INFO( @@ -966,6 +830,51 @@ 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; + } + 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 distinct NVRTC kernel configurations requested from the export + // library in this process (keyed as its kernel cache keys them). Grows only + // when a configuration is requested for the first time, i.e. on a compile, so + // a value that is stable across the first request after readiness proves no + // request-time JIT. + static int64_t numRequestedJITKernelKeys() + { + std::lock_guard interfaceLock(getFmhaInterfaceMutex()); + return static_cast(getRequestedJITKernelKeys().size()); + } + void run(RunnerParams const& params) { if (params.mMaxSeqLenQ == 0 || params.mBatchSize == 0 @@ -1075,6 +984,7 @@ class TllmGenFmhaKernel fmhaConfig.mGrid = grid; auto const compileStart = std::chrono::steady_clock::now(); mFmhaInterface.generateAndCompileKernel(fmhaConfig); + recordJITKernelKey(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 diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index 52c7ddc06526..0cc849741afb 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -52,6 +52,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..3bb23ec31759 --- /dev/null +++ b/cpp/tensorrt_llm/thop/trtllmGenFmhaJitWarmupOp.cpp @@ -0,0 +1,67 @@ +/* + * 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 distinct TRTLLM-Gen FMHA NVRTC kernel configurations requested in +// this process so far. A configuration is compiled the first time it is +// requested, so a value that does not change across a request proves that the +// request compiled nothing. +int64_t trtllmGenFmhaJitNumRequestedKernelKeys() +{ + return tensorrt_llm::kernels::TllmGenFmhaKernel::numRequestedJITKernelKeys(); +} + +// 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_requested_kernel_keys() -> int", + &tensorrt_llm::torch_ext::trtllmGenFmhaJitNumRequestedKernelKeys); + m.def("trtllm_gen_fmha_async_jit_warmup_enabled() -> bool", + &tensorrt_llm::torch_ext::trtllmGenFmhaAsyncJitWarmupEnabled); +} diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 313d0a6daf7b..7f9f77f2d264 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1579,14 +1579,7 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, # launch argument and is baked into every later replay. with self._warmup_timer.phase("cuda_graph_capture"), \ _moe_a2a_steady_state_budget_for_capture(): - with self.cuda_graph_runner.allow_capture(): - self.cuda_graph_runner.is_warmup_only = True - try: - self._run_cuda_graph_warmup(resource_manager) - finally: - self.cuda_graph_runner.is_warmup_only = False - self.cuda_graph_runner.padding_dummy_requests = {} - self._run_cuda_graph_warmup(resource_manager) + self._run_cuda_graph_warmup_and_capture(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") # Pre-compile DeepGEMM paged_mqa_logits_metadata for every 32-aligned # batch bucket the runtime can produce (max_batch_size scaled by the @@ -2643,6 +2636,47 @@ def _get_graphs_to_capture( return [(bs, draft_len) for bs in cuda_graph_batch_sizes for draft_len in draft_lengths] + def _run_cuda_graph_warmup_and_capture( + self, resource_manager: ResourceManager) -> None: + """Run the graph-shape warmup pass, then capture. + + The two passes are separated by the TRTLLM-Gen FMHA JIT warmup barrier: + every attention kernel the warmup derived must be compiled before any + launch is captured into a graph, and before the engine reports ready. + The barrier runs regardless of whether CUDA graphs or the general + warmup are enabled, so configurations without graphs still finish + verification before readiness. + """ + with self.cuda_graph_runner.allow_capture(): + self.cuda_graph_runner.is_warmup_only = True + try: + self._run_cuda_graph_warmup(resource_manager) + finally: + self.cuda_graph_runner.is_warmup_only = False + self.cuda_graph_runner.padding_dummy_requests = {} + self._drain_and_verify_trtllm_gen_fmha_jit_warmup() + self._run_cuda_graph_warmup(resource_manager) + + 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. A no-op unless async warmup is + enabled (``TRTLLM_GEN_FMHA_ASYNC_WARMUP=1``). + """ + if not issubclass(self.attn_backend.Metadata, TrtllmAttentionMetadata): + 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.") + def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): """Warm up or capture CUDA graphs for the configured graph shapes.""" is_warmup_only = self.cuda_graph_runner.is_warmup_only From 9bea9ea8823085942c7b11c249f4e37ab5f5e22e Mon Sep 17 00:00:00 2001 From: Aaron Meoded Date: Wed, 7 Oct 2026 15:08:40 -0400 Subject: [PATCH 06/11] [None][test] Cover the async FMHA JIT warmup: worker, registry, engine barrier, no post-ready JIT - cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp (host-only, registered in CMake): worker ordering, drain, failing tasks, shutdown with pending work, enqueue after shutdown; registry dedup, same-device multi-config isolation, multi-device isolation, claim-once semantics; and the scheduling protocol end to end including background failure retried on the foreground and inline fallback when the worker is shut down. - tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py (cpu_only, picked up by l0_cpu.yml's unittest/_torch/executor entry): the barrier runs after the graph-shape warmup pass and before capture, with CUDA graphs enabled and disabled, independently of the general-warmup gate, reports foreground recovery of failed background sweeps, is skipped for non-TRTLLM attention backends, and warmup() only reaches capture through the helper that carries it. - tests/unittest/_torch/attention/test_trtllm_gen_fmha_async_jit_warmup.py (Blackwell, picked up by l0_b200.yml's unittest/_torch/attention entry): boots Qwen3-0.6B with CUDA graphs in a fresh process per arm (async and sync warmup) and asserts capture succeeds, the warmup exercised the NVRTC path, and the first requests after readiness request no new kernel configuration, i.e. compile nothing. Signed-off-by: Aaron Meoded --- cpp/tests/unit_tests/kernels/CMakeLists.txt | 1 + .../unit_tests/kernels/fmhaJitWarmupTest.cpp | 413 ++++++++++++++++++ .../test_trtllm_gen_fmha_async_jit_warmup.py | 123 ++++++ .../executor/test_fmha_jit_warmup_barrier.py | 149 +++++++ 4 files changed, 686 insertions(+) create mode 100644 cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp create mode 100644 tests/unittest/_torch/attention/test_trtllm_gen_fmha_async_jit_warmup.py create mode 100644 tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index dc26b6f98ee0..c2f7eb630171 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -104,6 +104,7 @@ endif() add_gtest(sparseKvCacheTest sparseKvCacheTest.cu) add_gtest(prepareCustomMaskTest prepareCustomMaskTest.cpp) +add_gtest(fmhaJitWarmupTest fmhaJitWarmupTest.cpp) add_gtest(fmhaDsv4EpilogueFusionTest fmhaDsv4EpilogueFusionTest.cpp) set(NVFP4_COLD_PAGE_KERNEL_TEST_SRC diff --git a/cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp b/cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp new file mode 100644 index 000000000000..9a6b012a04e8 --- /dev/null +++ b/cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp @@ -0,0 +1,413 @@ +/* + * 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 and the sweep +// registry's state machine. 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::JITWarmupState; +using tensorrt_llm::kernels::TllmGenFmhaAsyncWarmupWorker; +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); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// +// Worker + registry together: the scheduling protocol used by fmhaKernels.h. +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// Mirrors runJITWarmupGridIfRequested(): register, queue the sweep, mark done on +// completion or failure, then verify on the foreground. +struct Scheduler +{ + TllmGenFmhaAsyncWarmupWorker mWorker; + Registry mRegistry; + std::atomic mBackgroundSweeps{0}; + std::atomic mForegroundSweeps{0}; + + void schedule(uint64_t fingerprint, FakeParams const& params, FakeKernel* kernel, bool failInBackground) + { + if (mRegistry.tryRegister(fingerprint, params, JITWarmupState::kScheduled, 0, kernel)) + { + return; + } + bool const queued = mWorker.enqueue( + [this, fingerprint, failInBackground]() + { + try + { + if (failInBackground) + { + throw std::runtime_error("nvrtc failed"); + } + ++mBackgroundSweeps; + } + catch (...) + { + mRegistry.markBackgroundDone(fingerprint); + throw; + } + mRegistry.markBackgroundDone(fingerprint); + }); + if (!queued) + { + ++mForegroundSweeps; + mRegistry.markBackgroundDone(fingerprint); + } + } + + int drainAndVerifyAll() + { + mWorker.drain(); + int numVerified = 0; + while (auto sweep = mRegistry.claimAnyUnverified()) + { + ++mForegroundSweeps; + ++numVerified; + } + return numVerified; + } +}; + +TEST(TllmGenFmhaJitWarmupProtocol, EveryScheduledSweepIsVerifiedExactlyOnceBeforeCapture) +{ + Scheduler scheduler; + FakeKernel kernel{1}; + scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, /*failInBackground=*/false); + scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, false); // duplicate: dropped + scheduler.schedule(0xB, FakeParams{1, 1024}, &kernel, false); + + 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}; + scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, /*failInBackground=*/true); + 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(); + scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, false); + // 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); +} + +} // namespace 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..1f528d09d730 --- /dev/null +++ b/tests/unittest/_torch/attention/test_trtllm_gen_fmha_async_jit_warmup.py @@ -0,0 +1,123 @@ +# 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] + keys = torch.ops.trtllm.trtllm_gen_fmha_jit_num_requested_kernel_keys + result = {"keys_before": keys()} + 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["keys_ready"] = keys() + 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["keys_after_first_requests"] = keys() + 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 = [l for l in proc.stdout.splitlines() if l.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"] + + # The warmup must have gone through the NVRTC path, otherwise the + # invariant below would hold vacuously. + assert result["keys_ready"] > result["keys_before"], ( + "warmup requested no NVRTC kernel configurations; this model/config " + "does not exercise the TRTLLM-Gen JIT path") + + # No configuration was requested for the first time after readiness, i.e. + # the first requests compiled nothing: every kernel they needed was + # compiled (and, for the async arm, verified) before capture. + assert result["keys_after_first_requests"] == result["keys_ready"], ( + f"{result['keys_after_first_requests'] - result['keys_ready']} kernel " + "configuration(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..a519d6953e9c --- /dev/null +++ b/tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py @@ -0,0 +1,149 @@ +# 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 engine warmup. + +The barrier (``torch.ops.trtllm.trtllm_gen_fmha_jit_warmup_drain_and_verify``) +must run after the graph-shape warmup pass and before any CUDA graph is +captured, and it must run whether or not CUDA graphs are enabled, so that +every attention kernel the warmup derived is compiled before readiness. These +tests drive the engine's capture sequence with the graph passes and the op +replaced by recorders; the C++ state machine behind the op is covered by +``cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp``. +""" + +import contextlib +import inspect +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 import model_engine as model_engine_module +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine +from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer + +pytestmark = pytest.mark.cpu_only + + +class _Recorder: + """Records the order of graph passes and barrier calls.""" + + def __init__(self, drain_result: int = 0): + self.events: list[str] = [] + self.drain_result = drain_result + + def run_cuda_graph_warmup(self, engine, resource_manager): + self.events.append( + "warmup" if engine.cuda_graph_runner.is_warmup_only else "capture") + + def drain_and_verify(self): + self.events.append("drain_and_verify") + return self.drain_result + + +def _engine(*, cuda_graphs_enabled: bool, + trtllm_attention: bool = True) -> PyTorchModelEngine: + """Build an engine carrying only what the capture sequence reads.""" + engine = object.__new__(PyTorchModelEngine) + + @contextlib.contextmanager + def allow_capture(): + yield + + engine.cuda_graph_runner = SimpleNamespace( + enabled=cuda_graphs_enabled, + is_warmup_only=False, + padding_dummy_requests={"stale": object()}, + allow_capture=allow_capture, + ) + engine.attn_backend = SimpleNamespace( + Metadata=TrtllmAttentionMetadata if trtllm_attention else + AttentionMetadata) + engine._warmup_timer = _WarmupTimer(rank=0) + return engine + + +@contextlib.contextmanager +def _patched(engine: PyTorchModelEngine, recorder: _Recorder): + with ( + mock.patch.object( + PyTorchModelEngine, + "_run_cuda_graph_warmup", + autospec=True, + side_effect=recorder.run_cuda_graph_warmup, + ), + mock.patch.object( + torch.ops.trtllm, + "trtllm_gen_fmha_jit_warmup_drain_and_verify", + create=True, + side_effect=recorder.drain_and_verify, + ), + ): + yield + + +@pytest.mark.parametrize("cuda_graphs_enabled", [True, False], + ids=["graphs_enabled", "graphs_disabled"]) +def test_barrier_runs_between_graph_warmup_and_capture( + cuda_graphs_enabled: bool) -> None: + engine = _engine(cuda_graphs_enabled=cuda_graphs_enabled) + recorder = _Recorder() + with _patched(engine, recorder): + engine._run_cuda_graph_warmup_and_capture(resource_manager=None) + + # The graph passes are invoked either way (they return early on their own + # when graphs are disabled); the barrier sits strictly between them, so + # verification finishes before the first captured launch. + assert recorder.events == ["warmup", "drain_and_verify", "capture"] + assert engine.cuda_graph_runner.is_warmup_only is False + assert engine.cuda_graph_runner.padding_dummy_requests == {} + + +def test_barrier_does_not_depend_on_general_warmup() -> None: + # The capture sequence runs the barrier unconditionally: it never consults + # the general-warmup gate (``can_run_general_warmup`` is decided earlier in + # ``warmup`` and only scopes the general / memory-pool phases). + engine = _engine(cuda_graphs_enabled=False) + recorder = _Recorder() + with _patched(engine, recorder): + engine._run_cuda_graph_warmup_and_capture(resource_manager=None) + assert recorder.events == ["warmup", "drain_and_verify", "capture"] + source = inspect.getsource( + PyTorchModelEngine._run_cuda_graph_warmup_and_capture) + assert "can_run_general_warmup" not in source + + +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 capture only after that. + engine = _engine(cuda_graphs_enabled=True) + recorder = _Recorder(drain_result=2) + with _patched(engine, recorder), \ + mock.patch.object(model_engine_module, "logger") as engine_logger: + engine._run_cuda_graph_warmup_and_capture(resource_manager=None) + assert recorder.events == ["warmup", "drain_and_verify", "capture"] + logged = " ".join( + str(arg) for call in engine_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: + engine = _engine(cuda_graphs_enabled=True, trtllm_attention=False) + recorder = _Recorder() + with _patched(engine, recorder): + engine._run_cuda_graph_warmup_and_capture(resource_manager=None) + assert recorder.events == ["warmup", "capture"] + + +def test_warmup_routes_capture_through_the_helper() -> None: + # ``warmup`` must not call ``_run_cuda_graph_warmup`` directly: the only + # path to capture is the helper that carries the barrier, so the barrier + # cannot be reordered or dropped by a change to ``warmup`` alone. + source = inspect.getsource(PyTorchModelEngine.warmup) + assert "self._run_cuda_graph_warmup_and_capture(" in source + assert "self._run_cuda_graph_warmup(" not in source From ea0455c4d6c16e595c7cc63d020f9d71f48bc0f0 Mon Sep 17 00:00:00 2001 From: Aaron Meoded Date: Wed, 7 Oct 2026 15:10:20 -0400 Subject: [PATCH 07/11] [None][chore] Apply clang-format Signed-off-by: Aaron Meoded --- cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 929934d42004..fad3c4410075 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -44,10 +44,10 @@ #include "tensorrt_llm/common/cudaDriverWrapper.h" #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h" #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" @@ -1491,7 +1491,7 @@ class TllmGenFmhaKernel fmhaData.mOutputBuffers.partialOPtrD = fmhaData.mOutputBuffers.partialStatsPtrD + partialStatsBufferSize; fmhaData.mOutputBuffers.skipSoftmaxStatsPtrD = nullptr; // Not available in params (would need to be added) fmhaData.mOutputBuffers.softmaxStatsD = params.softmaxStatsPtr; - fmhaData.mOutputBuffers.oDebugPtrD = nullptr; // Debug output not supported in TensorRT-LLM + fmhaData.mOutputBuffers.oDebugPtrD = nullptr; // Debug output not supported in TensorRT-LLM // Print all primitive type variables in FmhaData for debugging } From dc07ad4310da8eb98d00cd829e09a327561ede0f Mon Sep 17 00:00:00 2001 From: Aaron Meoded Date: Wed, 7 Oct 2026 16:25:25 -0400 Subject: [PATCH 08/11] [None][fix] Sync ranks after the JIT warmup drain; fix barrier test target; ruff - After the local drain-and-verify, ranks meet on the warmup agreement allgather before the capture forwards: each rank only waited for its own compilation, and the capture forwards run collectives under the steady-state MoE all-to-all budget. Both the drain and the sync are skipped when async warmup is disabled. - test_warmup_routes_capture_through_the_helper inspects _warmup_scheduled (where the capture call lives), not warmup. - ruff format / E741 on the two new test files. Signed-off-by: Aaron Meoded --- .../_torch/pyexecutor/model_engine.py | 17 ++- .../test_trtllm_gen_fmha_async_jit_warmup.py | 19 +-- .../executor/test_fmha_jit_warmup_barrier.py | 121 ++++++++++++------ 3 files changed, 108 insertions(+), 49 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 7f9f77f2d264..b6f3d2250638 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2669,13 +2669,22 @@ def _drain_and_verify_trtllm_gen_fmha_jit_warmup(self) -> None: """ if not issubclass(self.attn_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.") + 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 capture forwards + # below run collectives under the steady-state MoE all-to-all budget, + # so a rank that enters one while a peer is still compiling could + # exhaust that budget waiting. Meet the peers here first. + allgather = self._warmup_agreement_allgather() + if allgather is not None: + allgather(1) def _run_cuda_graph_warmup(self, resource_manager: ResourceManager): """Warm up or capture CUDA graphs for the configured graph shapes.""" 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 index 1f528d09d730..3716e5ee3608 100644 --- 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 @@ -88,10 +88,11 @@ def _run_arm(model_dir: Path, async_warmup: bool) -> dict: 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 = [l for l in proc.stdout.splitlines() if l.startswith("RESULT ")] + 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 "):]) + return json.loads(lines[-1][len("RESULT ") :]) @pytest.fixture(scope="module") @@ -100,10 +101,10 @@ def model_dir() -> Path: @skip_pre_blackwell -@pytest.mark.parametrize("async_warmup", [True, False], - ids=["async_warmup", "sync_warmup"]) +@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: + model_dir: Path, async_warmup: bool +) -> None: result = _run_arm(model_dir, async_warmup) assert result["async_enabled"] is async_warmup @@ -113,11 +114,13 @@ def test_first_requests_after_readiness_do_not_jit_compile( # invariant below would hold vacuously. assert result["keys_ready"] > result["keys_before"], ( "warmup requested no NVRTC kernel configurations; this model/config " - "does not exercise the TRTLLM-Gen JIT path") + "does not exercise the TRTLLM-Gen JIT path" + ) # No configuration was requested for the first time after readiness, i.e. # the first requests compiled nothing: every kernel they needed was # compiled (and, for the async arm, verified) before capture. assert result["keys_after_first_requests"] == result["keys_ready"], ( f"{result['keys_after_first_requests'] - result['keys_ready']} kernel " - "configuration(s) were compiled by the first requests after readiness") + "configuration(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 index a519d6953e9c..acd060af813c 100644 --- a/tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py +++ b/tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py @@ -29,23 +29,32 @@ class _Recorder: - """Records the order of graph passes and barrier calls.""" + """Records the order of graph passes, barrier calls and rank syncs.""" - def __init__(self, drain_result: int = 0): + 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 run_cuda_graph_warmup(self, engine, resource_manager): - self.events.append( - "warmup" if engine.cuda_graph_runner.is_warmup_only else "capture") + self.events.append("warmup" if engine.cuda_graph_runner.is_warmup_only else "capture") def drain_and_verify(self): self.events.append("drain_and_verify") return self.drain_result + def allgather(self, value): + self.events.append("rank_sync") + return [value, value] -def _engine(*, cuda_graphs_enabled: bool, - trtllm_attention: bool = True) -> PyTorchModelEngine: + +def _engine( + *, + cuda_graphs_enabled: bool, + trtllm_attention: bool = True, + recorder: _Recorder | None = None, + multi_rank: bool = False, +) -> PyTorchModelEngine: """Build an engine carrying only what the capture sequence reads.""" engine = object.__new__(PyTorchModelEngine) @@ -60,35 +69,46 @@ def allow_capture(): allow_capture=allow_capture, ) engine.attn_backend = SimpleNamespace( - Metadata=TrtllmAttentionMetadata if trtllm_attention else - AttentionMetadata) + Metadata=TrtllmAttentionMetadata if trtllm_attention else AttentionMetadata + ) engine._warmup_timer = _WarmupTimer(rank=0) + # The forward-group agreement primitive the barrier synchronizes on; None + # is what the real helper returns for a single rank (and under DWDP). + engine._warmup_agreement_allgather = ( + lambda: recorder.allgather if multi_rank and recorder else None + ) return engine @contextlib.contextmanager def _patched(engine: PyTorchModelEngine, recorder: _Recorder): with ( - mock.patch.object( - PyTorchModelEngine, - "_run_cuda_graph_warmup", - autospec=True, - side_effect=recorder.run_cuda_graph_warmup, - ), - 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( + PyTorchModelEngine, + "_run_cuda_graph_warmup", + autospec=True, + side_effect=recorder.run_cuda_graph_warmup, + ), + 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 -@pytest.mark.parametrize("cuda_graphs_enabled", [True, False], - ids=["graphs_enabled", "graphs_disabled"]) -def test_barrier_runs_between_graph_warmup_and_capture( - cuda_graphs_enabled: bool) -> None: +@pytest.mark.parametrize( + "cuda_graphs_enabled", [True, False], ids=["graphs_enabled", "graphs_disabled"] +) +def test_barrier_runs_between_graph_warmup_and_capture(cuda_graphs_enabled: bool) -> None: engine = _engine(cuda_graphs_enabled=cuda_graphs_enabled) recorder = _Recorder() with _patched(engine, recorder): @@ -111,8 +131,7 @@ def test_barrier_does_not_depend_on_general_warmup() -> None: with _patched(engine, recorder): engine._run_cuda_graph_warmup_and_capture(resource_manager=None) assert recorder.events == ["warmup", "drain_and_verify", "capture"] - source = inspect.getsource( - PyTorchModelEngine._run_cuda_graph_warmup_and_capture) + source = inspect.getsource(PyTorchModelEngine._run_cuda_graph_warmup_and_capture) assert "can_run_general_warmup" not in source @@ -122,13 +141,13 @@ def test_barrier_reports_foreground_recovery_of_background_failures() -> None: # to capture only after that. engine = _engine(cuda_graphs_enabled=True) recorder = _Recorder(drain_result=2) - with _patched(engine, recorder), \ - mock.patch.object(model_engine_module, "logger") as engine_logger: + with ( + _patched(engine, recorder), + mock.patch.object(model_engine_module, "logger") as engine_logger, + ): engine._run_cuda_graph_warmup_and_capture(resource_manager=None) assert recorder.events == ["warmup", "drain_and_verify", "capture"] - logged = " ".join( - str(arg) for call in engine_logger.info.call_args_list - for arg in call.args) + logged = " ".join(str(arg) for call in engine_logger.info.call_args_list for arg in call.args) assert "verified 2 background warmup sweep(s)" in logged @@ -140,10 +159,38 @@ def test_barrier_is_skipped_for_non_trtllm_attention_backends() -> None: assert recorder.events == ["warmup", "capture"] +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. + engine = _engine(cuda_graphs_enabled=True) + recorder = _Recorder(async_enabled=False) + with _patched(engine, recorder): + engine._run_cuda_graph_warmup_and_capture(resource_manager=None) + assert recorder.events == ["warmup", "capture"] + + +def test_ranks_meet_after_the_local_drain_and_before_capture() -> None: + # Each rank only waits for its own compilation; the capture forwards run + # collectives under the steady-state MoE all-to-all budget, so the ranks + # must meet after draining and before the first captured forward. + recorder = _Recorder() + engine = _engine(cuda_graphs_enabled=True, recorder=recorder, multi_rank=True) + with _patched(engine, recorder): + engine._run_cuda_graph_warmup_and_capture(resource_manager=None) + assert recorder.events == ["warmup", "drain_and_verify", "rank_sync", "capture"] + + def test_warmup_routes_capture_through_the_helper() -> None: - # ``warmup`` must not call ``_run_cuda_graph_warmup`` directly: the only - # path to capture is the helper that carries the barrier, so the barrier - # cannot be reordered or dropped by a change to ``warmup`` alone. - source = inspect.getsource(PyTorchModelEngine.warmup) - assert "self._run_cuda_graph_warmup_and_capture(" in source - assert "self._run_cuda_graph_warmup(" not in source + # The scheduled warmup must not call ``_run_cuda_graph_warmup`` directly: + # the only path to capture is the helper that carries the barrier, so the + # barrier cannot be reordered or dropped by a change to the warmup flow + # alone. ``warmup`` delegates to ``_warmup_impl`` and then + # ``_warmup_scheduled``; the capture call lives in the latter. + scheduled = inspect.getsource(inspect.unwrap(PyTorchModelEngine._warmup_scheduled)) + assert "self._run_cuda_graph_warmup_and_capture(" in scheduled + for method in ( + PyTorchModelEngine.warmup, + PyTorchModelEngine._warmup_impl, + PyTorchModelEngine._warmup_scheduled, + ): + assert "self._run_cuda_graph_warmup(" not in inspect.getsource(inspect.unwrap(method)) From f7b99d39befb127e7037c6ad0353849feeb6d9a3 Mon Sep 17 00:00:00 2001 From: Aaron Meoded Date: Wed, 7 Oct 2026 18:16:13 -0400 Subject: [PATCH 09/11] [None][fix] Count reported cache misses; make worker shutdown safe for every creation path Review feedback (chienchunhung): - The compile counter now tallies the KernelCacheStatus returned by both generateAndCompileKernel() call sites: misses (compiles, including the same key recompiled after an LRU eviction), hits, and Unknown. Two ops expose misses and unknown counts; the GPU test asserts unknown == 0 so an unreported status cannot satisfy it silently. Host test covers repeated misses for the same key and the unknown tally. - getExecPath() is now static and its string is leaked, so a sweep can read it at any point up to process exit regardless of which path created the worker. Worker creation is gated behind asyncJITWarmupWorkerCreated(): the pre-capture barrier, the capture-time drain, and kernel destruction no longer create a worker just to find it idle. With every sweep dependency leaked or load-time constructed and kernel objects draining in their destructor, the atexit shutdown no longer depends on static destruction order. - fmhaJitWarmupProcessExitTest.cpp (death test): empty barrier, worker created before the install path exists, work queued and still pending at exit(); the child exits 0 only if the pending work completed during exit and saw the same install path as the main thread. Signed-off-by: Aaron Meoded --- .../trtllmGenKernels/fmha/fmhaJitWarmup.h | 70 +++++++++-- .../trtllmGenKernels/fmha/fmhaKernels.h | 112 +++++++++++------- .../thop/trtllmGenFmhaJitWarmupOp.cpp | 23 ++-- cpp/tests/unit_tests/kernels/CMakeLists.txt | 1 + .../kernels/fmhaJitWarmupProcessExitTest.cpp | 91 ++++++++++++++ .../unit_tests/kernels/fmhaJitWarmupTest.cpp | 27 +++++ .../test_trtllm_gen_fmha_async_jit_warmup.py | 36 ++++-- 7 files changed, 288 insertions(+), 72 deletions(-) create mode 100644 cpp/tests/unit_tests/kernels/fmhaJitWarmupProcessExitTest.cpp diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h index f87b76e1772a..0fe97a3601da 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h @@ -49,16 +49,17 @@ namespace kernels // 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 through a -// std::atexit handler registered when the worker is first created, i.e. from -// the first warmup forward. atexit handlers run before the destruction of any -// static object whose initialization completed earlier, so a sweep still in -// flight at process exit finishes while everything it reads (the export -// library's globals, the kernel objects, the mutexes, the install-path string) -// is still alive. The one caveat is a static that the sweep itself constructs -// lazily on the worker thread: such a static would be destroyed before the -// handler runs. fmhaKernels.h therefore touches those on the registering -// thread before the first task is queued. +// 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 statics 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: @@ -336,6 +337,55 @@ class TllmGenFmhaJitWarmupRegistry std::atomic mNumUnverified{0}; }; +//////////////////////////////////////////////////////////////////////////////////////////////////// + +// 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 fad3c4410075..0c490feac5c0 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -160,8 +160,22 @@ inline bool asyncJITWarmupEnabled() 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. +// 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 = []() @@ -182,11 +196,10 @@ inline TllmGenFmhaAsyncWarmupWorker& getAsyncJITWarmupWorker() "run() will drain this worker before any CUDA graph capture."); } }); - // The worker is first created from the first warmup forward, after every - // static a sweep depends on has been constructed, so this handler drains - // and joins the worker before any of them is destroyed. The object itself - // is intentionally leaked. + // 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; @@ -268,12 +281,11 @@ class TllmGenFmhaKernel ~TllmGenFmhaKernel() { - // A queued background warmup task captures `this`. At process exit the - // worker's atexit handler has already drained and joined it (see - // fmhaJitWarmup.h); draining here as well covers a kernel object destroyed - // while the process keeps running. The worker is leaked, so this is safe - // during static destruction too. - if (asyncJITWarmupEnabled()) + // 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(); } @@ -541,8 +553,7 @@ class TllmGenFmhaKernel fmhaConfig.mCtaDim = ctaDim; fmhaConfig.mGrid = grid; auto const compileStart = std::chrono::steady_clock::now(); - mFmhaInterface.generateAndCompileKernel(fmhaConfig); - recordJITKernelKey(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 @@ -613,20 +624,27 @@ class TllmGenFmhaKernel return *sRegistry; } - // Keys of every NVRTC kernel configuration requested from the export library - // so far, hashed exactly as its kernel cache keys them: a key seen for the - // first time is a compile. Intentionally leaked. Guarded by - // getFmhaInterfaceMutex(), which both request sites already hold. - static std::unordered_set& getRequestedJITKernelKeys() + // Tally of the export library's kernel-cache results (one per + // generateAndCompileKernel() call). Intentionally leaked. + static TllmGenFmhaJitCompileStats& getJITCompileStats() { - static auto* const sKeys = new std::unordered_set(); - return *sKeys; + static auto* const sStats = new TllmGenFmhaJitCompileStats(); + return *sStats; } - // Must be called with getFmhaInterfaceMutex() held, after generateAndCompileKernel(). - static void recordJITKernelKey(FmhaConfig const& fmhaConfig) + static void recordJITCacheResult(FmhaInterface::KernelCacheStatus status) { - getRequestedJITKernelKeys().insert(std::hash{}(fmhaConfig.mOptions)); + 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); } // Registers a warmup sweep for this kernel object on the current device; @@ -690,7 +708,8 @@ class TllmGenFmhaKernel // populated kernel cache. void maybeWaitForAsyncJITWarmup(cudaStream_t stream) const { - if (!asyncJITWarmupEnabled() || getAsyncJITWarmupWorker().numPendingTasks() == 0) + if (!asyncJITWarmupEnabled() || !asyncJITWarmupWorkerCreated().load(std::memory_order_acquire) + || getAsyncJITWarmupWorker().numPendingTasks() == 0) { return; } @@ -741,10 +760,6 @@ class TllmGenFmhaKernel if (asyncJITWarmupEnabled()) { - // The install-path string is a function-local static that the sweep reads - // on the worker thread. Construct it on this thread, before the worker and - // its atexit shutdown exist (lifetime note in fmhaJitWarmup.h). - (void) getExecPath(); uint64_t const fingerprint = jitWarmupFingerprint(warmupParams); int const deviceId = tensorrt_llm::common::getDevice(); TLLM_LOG_INFO( @@ -842,7 +857,10 @@ class TllmGenFmhaKernel { return 0; } - getAsyncJITWarmupWorker().drain(); + if (asyncJITWarmupWorkerCreated().load(std::memory_order_acquire)) + { + getAsyncJITWarmupWorker().drain(); + } struct RestoreDevice { @@ -864,15 +882,21 @@ class TllmGenFmhaKernel return numVerified; } - // Number of distinct NVRTC kernel configurations requested from the export - // library in this process (keyed as its kernel cache keys them). Grows only - // when a configuration is requested for the first time, i.e. on a compile, so - // a value that is stable across the first request after readiness proves no - // request-time JIT. - static int64_t numRequestedJITKernelKeys() + // 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() { - std::lock_guard interfaceLock(getFmhaInterfaceMutex()); - return static_cast(getRequestedJITKernelKeys().size()); + return getJITCompileStats().numUnknown(); } void run(RunnerParams const& params) @@ -983,8 +1007,7 @@ class TllmGenFmhaKernel fmhaConfig.mCtaDim = ctaDim; fmhaConfig.mGrid = grid; auto const compileStart = std::chrono::steady_clock::now(); - mFmhaInterface.generateAndCompileKernel(fmhaConfig); - recordJITKernelKey(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 @@ -1255,9 +1278,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 @@ -1388,6 +1417,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/trtllmGenFmhaJitWarmupOp.cpp b/cpp/tensorrt_llm/thop/trtllmGenFmhaJitWarmupOp.cpp index 3bb23ec31759..0b8bb41f8d04 100644 --- a/cpp/tensorrt_llm/thop/trtllmGenFmhaJitWarmupOp.cpp +++ b/cpp/tensorrt_llm/thop/trtllmGenFmhaJitWarmupOp.cpp @@ -36,13 +36,19 @@ int64_t trtllmGenFmhaJitWarmupDrainAndVerify() return tensorrt_llm::kernels::TllmGenFmhaKernel::drainAndVerifyAllJITWarmups(); } -// Number of distinct TRTLLM-Gen FMHA NVRTC kernel configurations requested in -// this process so far. A configuration is compiled the first time it is -// requested, so a value that does not change across a request proves that the -// request compiled nothing. -int64_t trtllmGenFmhaJitNumRequestedKernelKeys() +// 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::numRequestedJITKernelKeys(); + 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 @@ -60,8 +66,9 @@ 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_requested_kernel_keys() -> int", - &tensorrt_llm::torch_ext::trtllmGenFmhaJitNumRequestedKernelKeys); + 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 c2f7eb630171..b52142c37112 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -105,6 +105,7 @@ endif() add_gtest(sparseKvCacheTest sparseKvCacheTest.cu) add_gtest(prepareCustomMaskTest prepareCustomMaskTest.cpp) add_gtest(fmhaJitWarmupTest fmhaJitWarmupTest.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/fmhaJitWarmupProcessExitTest.cpp b/cpp/tests/unit_tests/kernels/fmhaJitWarmupProcessExitTest.cpp new file mode 100644 index 000000000000..6a64748824d3 --- /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 statics 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 index 9a6b012a04e8..0beb0b86c36e 100644 --- a/cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp +++ b/cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp @@ -36,6 +36,7 @@ namespace using tensorrt_llm::kernels::JITWarmupState; 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. @@ -309,6 +310,32 @@ TEST(TllmGenFmhaJitWarmupRegistry, ClaimAnyUnverifiedServesTheDrainAndVerifyBarr 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. //////////////////////////////////////////////////////////////////////////////////////////////////// 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 index 3716e5ee3608..42538ec55129 100644 --- 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 @@ -42,8 +42,9 @@ from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig model_dir = sys.argv[1] - keys = torch.ops.trtllm.trtllm_gen_fmha_jit_num_requested_kernel_keys - result = {"keys_before": keys()} + 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", @@ -56,7 +57,7 @@ with llm: result["async_enabled"] = bool( torch.ops.trtllm.trtllm_gen_fmha_async_jit_warmup_enabled()) - result["keys_ready"] = keys() + result["misses_ready"] = misses() prompts = [ "The capital of France is", "Write one sentence about the ocean.", @@ -67,7 +68,8 @@ result["num_outputs"] = len(outputs) result["all_generated"] = all( len(o.outputs[0].token_ids) > 0 for o in outputs) - result["keys_after_first_requests"] = keys() + result["misses_after_first_requests"] = misses() + result["unknown_cache_results"] = unknown() print("RESULT " + json.dumps(result)) """) @@ -110,17 +112,25 @@ def test_first_requests_after_readiness_do_not_jit_compile( 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["keys_ready"] > result["keys_before"], ( - "warmup requested no NVRTC kernel configurations; this model/config " - "does not exercise the TRTLLM-Gen JIT path" + 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 configuration was requested for the first time after readiness, i.e. - # the first requests compiled nothing: every kernel they needed was - # compiled (and, for the async arm, verified) before capture. - assert result["keys_after_first_requests"] == result["keys_ready"], ( - f"{result['keys_after_first_requests'] - result['keys_ready']} kernel " - "configuration(s) were compiled by the first requests after readiness" + # 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" ) From e1280c2a45d29eda25238bf3e841d2ccee816b68 Mon Sep 17 00:00:00 2001 From: Aaron Meoded Date: Fri, 9 Oct 2026 17:51:47 -0400 Subject: [PATCH 10/11] [None][chore] Pre-commit: restore upstream comment alignment, codespell Signed-off-by: Aaron Meoded --- cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h | 2 +- cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h | 2 +- cpp/tests/unit_tests/kernels/fmhaJitWarmupProcessExitTest.cpp | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h index 0fe97a3601da..8d43e30475ce 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h @@ -52,7 +52,7 @@ namespace kernels // 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 statics is not relied upon: every +// 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), diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 0c490feac5c0..65efcc2ae8c6 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -1521,7 +1521,7 @@ class TllmGenFmhaKernel fmhaData.mOutputBuffers.partialOPtrD = fmhaData.mOutputBuffers.partialStatsPtrD + partialStatsBufferSize; fmhaData.mOutputBuffers.skipSoftmaxStatsPtrD = nullptr; // Not available in params (would need to be added) fmhaData.mOutputBuffers.softmaxStatsD = params.softmaxStatsPtr; - fmhaData.mOutputBuffers.oDebugPtrD = nullptr; // Debug output not supported in TensorRT-LLM + fmhaData.mOutputBuffers.oDebugPtrD = nullptr; // Debug output not supported in TensorRT-LLM // Print all primitive type variables in FmhaData for debugging } diff --git a/cpp/tests/unit_tests/kernels/fmhaJitWarmupProcessExitTest.cpp b/cpp/tests/unit_tests/kernels/fmhaJitWarmupProcessExitTest.cpp index 6a64748824d3..146de3b6e232 100644 --- a/cpp/tests/unit_tests/kernels/fmhaJitWarmupProcessExitTest.cpp +++ b/cpp/tests/unit_tests/kernels/fmhaJitWarmupProcessExitTest.cpp @@ -58,7 +58,7 @@ std::string gMainThreadExecPath; // "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 statics that do not exist yet. + // depends on static objects that do not exist yet. EXPECT_EQ(TllmGenFmhaKernel::drainAndVerifyAllJITWarmups(), 0); EXPECT_FALSE(asyncJITWarmupWorkerCreated().load()); From 98d97bbc79b1779e0fd535655a1a12a7e93169fa Mon Sep 17 00:00:00 2001 From: Aaron Meoded Date: Wed, 14 Oct 2026 23:09:30 -0400 Subject: [PATCH 11/11] [None][fix] Run the JIT warmup barrier before the steady-state A2A budget; hash DSv4 fusion; share the scheduling path Review feedback (lori-ren): - The drain-and-verify barrier and the rank sync now run before moe_a2a_steady_state_budget_for_capture() is entered, ahead of both the graph-shape warmup pass and capture. Every sweep is queued during the attention warmup, and the graph-shape warmup pass runs eager forwards with real collectives under the steady-state budget, so a compile (or a wait on the interface mutex behind a background compile) between two all-to-alls on one rank would spend it. The budget switch and the capture-timer phase move into _run_cuda_graph_warmup_and_capture() so the ordering is owned and tested in one place; the drain phase is no longer nested in the recorded capture phase. - jitWarmupFingerprint() hashes params.mDsv4EpilogueFusion.enabled, which selectKernel() reads (fused epilogue, forced E4M3 output), and takes the device as an argument so it is a pure function of its inputs. - The register-then-queue-or-run-inline glue is lifted out of runJITWarmupGridIfRequested() into scheduleJITWarmupSweep() in fmhaJitWarmup.h; the kernel and the protocol tests call the same code. The test Scheduler is now a thin wrapper, with new cases for sync mode (inline, still deduplicating) and for the sweep receiving its own device. - fmhaJitWarmupFingerprintTest.cpp covers determinism, pointer/stream independence, and separation by device, kernel-object dtype, spec-decoding and grid params, and the DSv4 fusion flag. - test_fmha_jit_warmup_barrier.py asserts the barrier precedes the budget switch and both graph passes, that the timer phases are siblings, and drives _warmup_scheduled() itself (device-bound phases stubbed) to the barrier; the duplicate general-warmup test and the source-text checks are gone. Signed-off-by: Aaron Meoded --- .../trtllmGenKernels/fmha/fmhaJitWarmup.h | 68 ++++++++ .../trtllmGenKernels/fmha/fmhaKernels.h | 124 ++++++--------- cpp/tests/unit_tests/kernels/CMakeLists.txt | 1 + .../kernels/fmhaJitWarmupFingerprintTest.cpp | 145 ++++++++++++++++++ .../unit_tests/kernels/fmhaJitWarmupTest.cpp | 100 ++++++++---- .../engine/runners/decoder/runner.py | 51 +++--- .../executor/test_fmha_jit_warmup_barrier.py | 128 +++++++++++----- 7 files changed, 439 insertions(+), 178 deletions(-) create mode 100644 cpp/tests/unit_tests/kernels/fmhaJitWarmupFingerprintTest.cpp diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h index 8d43e30475ce..1ae3691660c8 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaJitWarmup.h @@ -339,6 +339,74 @@ class TllmGenFmhaJitWarmupRegistry //////////////////////////////////////////////////////////////////////////////////////////////////// +// 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 diff --git a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h index 65efcc2ae8c6..3db8f6059035 100644 --- a/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h +++ b/cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h @@ -564,14 +564,19 @@ class TllmGenFmhaKernel } } +public: // The set of kernels a warmup sweep compiles only depends on this kernel - // instance (dtypes/SM/sage), the current device (modules load into the - // calling device's CUDA context), 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. - uint64_t jitWarmupFingerprint(RunnerParams const& params) const + // 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)); }; @@ -582,23 +587,18 @@ class TllmGenFmhaKernel uint32_t scaleQBits = 0; std::memcpy(&scaleQBits, ¶ms.mScaleQ, sizeof(scaleQBits)); uint64_t fingerprint = 0; - // Include the current device: compiled modules are loaded into the calling - // device's CUDA context, so an identical configuration on another device - // still needs its own sweep. Warmup requests always run with the target - // device current. for (uint64_t value : - std::initializer_list{static_cast(tensorrt_llm::common::getDevice()), - 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), + 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), @@ -607,13 +607,15 @@ class TllmGenFmhaKernel 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.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 @@ -647,19 +649,6 @@ class TllmGenFmhaKernel getJITCompileStats().record(result); } - // Registers a warmup sweep for this kernel object on the current device; - // returns true if an identical sweep was already registered. - bool tryRegisterJITWarmupConfig(RunnerParams const& params, JITWarmupState initialState) - { - return getJITWarmupRegistry().tryRegister( - jitWarmupFingerprint(params), params, initialState, tensorrt_llm::common::getDevice(), this); - } - - static void markJITWarmupBackgroundDone(uint64_t fingerprint) - { - getJITWarmupRegistry().markBackgroundDone(fingerprint); - } - // 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 @@ -745,59 +734,32 @@ class TllmGenFmhaKernel TLLM_CHECK_WITH_INFO(maxBatchSize > 0 && maxSeqLenKv > 0 && (!useGenKernelForPrefill || maxSeqLenQ > 0), "TRTLLM-Gen Fmha Warmup Param is invalid."); - // Deduplicate identical sweeps: 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). // 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; - if (tryRegisterJITWarmupConfig( - warmupParams, asyncJITWarmupEnabled() ? JITWarmupState::kScheduled : JITWarmupState::kVerified)) - { - return; - } - - if (asyncJITWarmupEnabled()) + 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) { - uint64_t const fingerprint = jitWarmupFingerprint(warmupParams); - int const deviceId = tensorrt_llm::common::getDevice(); TLLM_LOG_INFO( - "TRTLLM-Gen FMHA JIT warmup: scheduling background compilation of the generation-kernel grid " + "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)); - bool const queued = getAsyncJITWarmupWorker().enqueue( - [this, warmupParams, deviceId, fingerprint]() - { - try - { - TLLM_CUDA_CHECK(cudaSetDevice(deviceId)); - runJITWarmupGrid(warmupParams); - } - catch (...) - { - // Mark done even on failure so the fingerprint is not stranded in - // kScheduled: the foreground verify pass then re-runs this sweep - // and compiles whatever the background pass did not finish, - // instead of dedupe silently suppressing it forever. - markJITWarmupBackgroundDone(fingerprint); - throw; - } - markJITWarmupBackgroundDone(fingerprint); - TLLM_LOG_INFO( - "TRTLLM-Gen FMHA JIT warmup: background compilation of the generation-kernel grid " - "finished (fingerprint=%llx).", - static_cast(fingerprint)); - }); - if (!queued) - { - runJITWarmupGrid(runnerParams); - markJITWarmupBackgroundDone(fingerprint); - } - return; } - - runJITWarmupGrid(runnerParams); } void runJITWarmupGrid(RunnerParams const& runnerParams) diff --git a/cpp/tests/unit_tests/kernels/CMakeLists.txt b/cpp/tests/unit_tests/kernels/CMakeLists.txt index b52142c37112..bbb942ee6b88 100644 --- a/cpp/tests/unit_tests/kernels/CMakeLists.txt +++ b/cpp/tests/unit_tests/kernels/CMakeLists.txt @@ -105,6 +105,7 @@ 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) 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/fmhaJitWarmupTest.cpp b/cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp index 0beb0b86c36e..696c4302c72d 100644 --- a/cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp +++ b/cpp/tests/unit_tests/kernels/fmhaJitWarmupTest.cpp @@ -15,9 +15,9 @@ */ // Host-only tests for the asynchronous TRTLLM-Gen FMHA JIT warmup building -// blocks (fmhaJitWarmup.h): the background worker's lifecycle and the sweep -// registry's state machine. No GPU is involved; the "sweeps" here are plain -// callbacks. +// 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" @@ -34,7 +34,9 @@ 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; @@ -340,46 +342,40 @@ TEST(TllmGenFmhaJitCompileStats, CountsEveryMissIncludingRepeatsOfTheSameKey) // Worker + registry together: the scheduling protocol used by fmhaKernels.h. //////////////////////////////////////////////////////////////////////////////////////////////////// -// Mirrors runJITWarmupGridIfRequested(): register, queue the sweep, mark done on -// completion or failure, then verify on the foreground. +// 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 { - TllmGenFmhaAsyncWarmupWorker mWorker; + // 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}; - void schedule(uint64_t fingerprint, FakeParams const& params, FakeKernel* kernel, bool failInBackground) + JITWarmupScheduleResult schedule(uint64_t fingerprint, FakeParams const& params, FakeKernel* kernel, + bool failInBackground, bool async = true, int deviceId = 0) { - if (mRegistry.tryRegister(fingerprint, params, JITWarmupState::kScheduled, 0, kernel)) - { - return; - } - bool const queued = mWorker.enqueue( - [this, fingerprint, failInBackground]() + return scheduleJITWarmupSweep(mRegistry, async ? &mWorker : nullptr, fingerprint, params, deviceId, kernel, + [this, failInBackground](FakeParams const&, int) { - try + if (std::this_thread::get_id() == mMainThread) { - if (failInBackground) - { - throw std::runtime_error("nvrtc failed"); - } - ++mBackgroundSweeps; + ++mForegroundSweeps; + return; } - catch (...) + if (failInBackground) { - mRegistry.markBackgroundDone(fingerprint); - throw; + throw std::runtime_error("nvrtc failed"); } - mRegistry.markBackgroundDone(fingerprint); + ++mBackgroundSweeps; }); - if (!queued) - { - ++mForegroundSweeps; - mRegistry.markBackgroundDone(fingerprint); - } } + // The pre-capture barrier: drain, then re-run every finished sweep once. int drainAndVerifyAll() { mWorker.drain(); @@ -397,9 +393,10 @@ TEST(TllmGenFmhaJitWarmupProtocol, EveryScheduledSweepIsVerifiedExactlyOnceBefor { Scheduler scheduler; FakeKernel kernel{1}; - scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, /*failInBackground=*/false); - scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, false); // duplicate: dropped - scheduler.schedule(0xB, FakeParams{1, 1024}, &kernel, false); + 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); @@ -413,7 +410,8 @@ TEST(TllmGenFmhaJitWarmupProtocol, BackgroundFailureIsRetriedOnTheForeground) { Scheduler scheduler; FakeKernel kernel{1}; - scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, /*failInBackground=*/true); + 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); @@ -429,7 +427,7 @@ TEST(TllmGenFmhaJitWarmupProtocol, ShutDownWorkerFallsBackToInlineCompilation) Scheduler scheduler; FakeKernel kernel{1}; scheduler.mWorker.shutdown(); - scheduler.schedule(0xA, FakeParams{4, 1024}, &kernel, false); + 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); @@ -437,4 +435,40 @@ TEST(TllmGenFmhaJitWarmupProtocol, ShutDownWorkerFallsBackToInlineCompilation) 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/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 29d8663d44dc..2cc7c0450112 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -773,18 +773,7 @@ def _warmup_scheduled(self, resource_manager: ResourceManager, kv_cache_manager) # NVSHMEM). gc.collect() torch.cuda.empty_cache() - # Warm up every graph shape before capturing any graph. Attention - # kernels can switch implementations at smaller batch sizes and require - # a larger workspace, so the first pass grows the workspace to its - # maximum size. The second pass runs the final per-shape warmup and - # captures without resizing the workspace. - # Capture with the steady-state MoE all-to-all budget: the timeout is a - # launch argument and is baked into every later replay. - with ( - self._warmup_timer.phase("cuda_graph_capture"), - moe_a2a_steady_state_budget_for_capture(), - ): - self._run_cuda_graph_warmup_and_capture(resource_manager) + self._run_cuda_graph_warmup_and_capture(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") # Pre-compile DeepGEMM paged_mqa_logits_metadata for every 32-aligned # batch bucket the runtime can produce (max_batch_size scaled by the @@ -1902,23 +1891,36 @@ def _get_graphs_to_capture(self, cuda_graph_batch_sizes: list[int]) -> list[tupl return [(bs, draft_len) for bs in cuda_graph_batch_sizes for draft_len in draft_lengths] def _run_cuda_graph_warmup_and_capture(self, resource_manager: ResourceManager) -> None: - """Run the graph-shape warmup pass, then capture. - - The two passes are separated by the TRTLLM-Gen FMHA JIT warmup barrier: - every attention kernel the warmup derived must be compiled before any - launch is captured into a graph, and before the engine reports ready. - The barrier runs regardless of whether CUDA graphs or the general + """Finish the TRTLLM-Gen FMHA JIT warmup, then warm up every graph shape and capture. + + The JIT warmup barrier comes first, before the steady-state MoE + all-to-all budget is entered: every sweep was queued during the + attention warmup, and the graph-shape warmup pass runs eager forwards + with real collectives under that budget, so compiling (or waiting on a + background compile) between two all-to-alls on one rank would spend + it. The barrier runs regardless of whether CUDA graphs or the general warmup are enabled, so configurations without graphs still finish verification before readiness. + + Attention kernels can switch implementations at smaller batch sizes + and require a larger workspace, so the first graph pass grows the + workspace to its maximum size. The second pass runs the final + per-shape warmup and captures without resizing the workspace. Both + run with the steady-state budget: the timeout is a launch argument + and is baked into every later replay. """ - with self.cuda_graph_runner.allow_capture(): + self._drain_and_verify_trtllm_gen_fmha_jit_warmup() + with ( + self._warmup_timer.phase("cuda_graph_capture"), + moe_a2a_steady_state_budget_for_capture(), + self.cuda_graph_runner.allow_capture(), + ): self.cuda_graph_runner.is_warmup_only = True try: self._run_cuda_graph_warmup(resource_manager) finally: self.cuda_graph_runner.is_warmup_only = False self.cuda_graph_runner.padding_dummy_requests = {} - self._drain_and_verify_trtllm_gen_fmha_jit_warmup() self._run_cuda_graph_warmup(resource_manager) def _drain_and_verify_trtllm_gen_fmha_jit_warmup(self) -> None: @@ -1942,10 +1944,11 @@ def _drain_and_verify_trtllm_gen_fmha_jit_warmup(self) -> None: 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 capture forwards - # below run collectives under the steady-state MoE all-to-all budget, - # so a rank that enters one while a peer is still compiling could - # exhaust that budget waiting. Meet the peers here first. + # Each rank only waited for its own compilation. The graph-shape + # warmup and capture forwards that follow run collectives under the + # steady-state MoE all-to-all budget, so a rank that enters one + # while a peer is still compiling could exhaust that budget + # waiting. Meet the peers here first. allgather = self._warmup_agreement_allgather() if allgather is not None: allgather(1) diff --git a/tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py b/tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py index 644e75b8bf0f..7dba05e38f4f 100644 --- a/tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py +++ b/tests/unittest/_torch/executor/test_fmha_jit_warmup_barrier.py @@ -2,17 +2,18 @@ # SPDX-License-Identifier: Apache-2.0 """Ordering tests for the TRTLLM-Gen FMHA JIT warmup barrier in engine warmup. -The barrier (``torch.ops.trtllm.trtllm_gen_fmha_jit_warmup_drain_and_verify``) -must run after the graph-shape warmup pass and before any CUDA graph is -captured, and it must run whether or not CUDA graphs are enabled, so that -every attention kernel the warmup derived is compiled before readiness. These -tests drive the runner's capture sequence with the graph passes and the ops -replaced by recorders; the C++ state machine behind the ops is covered by +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 steady-state MoE all-to-all budget is entered for the graph-shape warmup +and capture passes, and it must run whether or not CUDA graphs are enabled, +so that every attention kernel the warmup derived is compiled before +readiness. These tests drive the runner's capture sequence with the graph +passes, the budget 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 -import inspect from types import SimpleNamespace from unittest import mock @@ -29,7 +30,7 @@ class _Recorder: - """Records the order of graph passes, barrier calls and rank syncs.""" + """Records the order of graph passes, barrier calls, rank syncs and budget switches.""" def __init__(self, drain_result: int = 0, async_enabled: bool = True): self.events: list[str] = [] @@ -47,6 +48,14 @@ def allgather(self, value): self.events.append("rank_sync") return [value, value] + @contextlib.contextmanager + def steady_state_budget(self): + self.events.append("budget_enter") + try: + yield + finally: + self.events.append("budget_exit") + def _runner( *, @@ -67,6 +76,7 @@ def allow_capture(): is_warmup_only=False, padding_dummy_requests={"stale": object()}, allow_capture=allow_capture, + preallocate_padding_dummies=mock.Mock(), ) runner._config = SimpleNamespace( attention_backend=SimpleNamespace( @@ -74,6 +84,7 @@ def allow_capture(): ) ) runner._warmup_timer = _WarmupTimer(rank=0) + runner._metrics = {} # The forward-group agreement primitive the barrier synchronizes on; None # is what the real helper returns for a single rank (and under DWDP). runner._warmup_agreement_allgather = ( @@ -91,6 +102,11 @@ def _patched(recorder: _Recorder): autospec=True, side_effect=recorder.run_cuda_graph_warmup, ), + mock.patch.object( + runner_module, + "moe_a2a_steady_state_budget_for_capture", + side_effect=recorder.steady_state_budget, + ), mock.patch.object( torch.ops.trtllm, "trtllm_gen_fmha_jit_warmup_drain_and_verify", @@ -107,34 +123,37 @@ def _patched(recorder: _Recorder): yield +_CAPTURE_UNDER_BUDGET = ["budget_enter", "warmup", "capture", "budget_exit"] + + @pytest.mark.parametrize( "cuda_graphs_enabled", [True, False], ids=["graphs_enabled", "graphs_disabled"] ) -def test_barrier_runs_between_graph_warmup_and_capture(cuda_graphs_enabled: bool) -> None: +def test_barrier_runs_before_the_steady_state_budget(cuda_graphs_enabled: bool) -> None: runner = _runner(cuda_graphs_enabled=cuda_graphs_enabled) recorder = _Recorder() with _patched(recorder): runner._run_cuda_graph_warmup_and_capture(resource_manager=None) # The graph passes are invoked either way (they return early on their own - # when graphs are disabled); the barrier sits strictly between them, so - # verification finishes before the first captured launch. - assert recorder.events == ["warmup", "drain_and_verify", "capture"] + # when graphs are disabled). The barrier precedes the budget switch, so no + # forward runs under the steady-state all-to-all budget while a compile + # may still be in flight. + assert recorder.events == ["drain_and_verify", *_CAPTURE_UNDER_BUDGET] assert runner.cuda_graph_runner.is_warmup_only is False assert runner.cuda_graph_runner.padding_dummy_requests == {} -def test_barrier_does_not_depend_on_general_warmup() -> None: - # The capture sequence runs the barrier unconditionally: it never consults - # the general-warmup gate (``can_run_general_warmup`` is decided earlier in - # the scheduled warmup and only scopes the general / memory-pool phases). - runner = _runner(cuda_graphs_enabled=False) +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). + runner = _runner(cuda_graphs_enabled=True) recorder = _Recorder() with _patched(recorder): runner._run_cuda_graph_warmup_and_capture(resource_manager=None) - assert recorder.events == ["warmup", "drain_and_verify", "capture"] - source = inspect.getsource(DecoderRunner._run_cuda_graph_warmup_and_capture) - assert "can_run_general_warmup" not in source + timings = runner._warmup_timer.timings + assert set(timings) == {"attention_jit_drain_and_verify", "cuda_graph_capture"} + assert recorder.events.index("drain_and_verify") < recorder.events.index("budget_enter") def test_barrier_reports_foreground_recovery_of_background_failures() -> None: @@ -145,7 +164,7 @@ def test_barrier_reports_foreground_recovery_of_background_failures() -> None: recorder = _Recorder(drain_result=2) with _patched(recorder), mock.patch.object(runner_module, "logger") as runner_logger: runner._run_cuda_graph_warmup_and_capture(resource_manager=None) - assert recorder.events == ["warmup", "drain_and_verify", "capture"] + assert recorder.events == ["drain_and_verify", *_CAPTURE_UNDER_BUDGET] logged = " ".join(str(arg) for call in runner_logger.info.call_args_list for arg in call.args) assert "verified 2 background warmup sweep(s)" in logged @@ -155,7 +174,7 @@ def test_barrier_is_skipped_for_non_trtllm_attention_backends() -> None: recorder = _Recorder() with _patched(recorder): runner._run_cuda_graph_warmup_and_capture(resource_manager=None) - assert recorder.events == ["warmup", "capture"] + assert recorder.events == _CAPTURE_UNDER_BUDGET def test_barrier_is_skipped_when_async_warmup_is_disabled() -> None: @@ -165,27 +184,56 @@ def test_barrier_is_skipped_when_async_warmup_is_disabled() -> None: recorder = _Recorder(async_enabled=False) with _patched(recorder): runner._run_cuda_graph_warmup_and_capture(resource_manager=None) - assert recorder.events == ["warmup", "capture"] + assert recorder.events == _CAPTURE_UNDER_BUDGET -def test_ranks_meet_after_the_local_drain_and_before_capture() -> None: - # Each rank only waits for its own compilation; the capture forwards run - # collectives under the steady-state MoE all-to-all budget, so the ranks - # must meet after draining and before the first captured forward. +def test_ranks_meet_after_the_local_drain_and_before_the_budget() -> None: + # Each rank only waits for its own compilation; the forwards under the + # steady-state MoE all-to-all budget would otherwise charge a peer's + # remaining compile time against that budget, so the ranks meet after + # draining and before the budget is entered. recorder = _Recorder() runner = _runner(cuda_graphs_enabled=True, recorder=recorder, multi_rank=True) with _patched(recorder): runner._run_cuda_graph_warmup_and_capture(resource_manager=None) - assert recorder.events == ["warmup", "drain_and_verify", "rank_sync", "capture"] - - -def test_warmup_routes_capture_through_the_helper() -> None: - # The scheduled warmup must not call ``_run_cuda_graph_warmup`` directly: - # the only path to capture is the helper that carries the barrier, so the - # barrier cannot be reordered or dropped by a change to the warmup flow - # alone. ``warmup`` delegates to ``_warmup_scheduled``, where the capture - # call lives. - scheduled = inspect.getsource(inspect.unwrap(DecoderRunner._warmup_scheduled)) - assert "self._run_cuda_graph_warmup_and_capture(" in scheduled - for method in (DecoderRunner.warmup, DecoderRunner._warmup_scheduled): - assert "self._run_cuda_graph_warmup(" not in inspect.getsource(inspect.unwrap(method)) + assert recorder.events == ["drain_and_verify", "rank_sync", *_CAPTURE_UNDER_BUDGET] + + +def test_scheduled_warmup_reaches_the_barrier_before_capture() -> None: + # Drive the scheduled warmup itself, with every phase that needs a device + # or a model stubbed out, and check that it arrives at the barrier through + # the capture helper: the drain happens, and it happens before the budget + # switch and the graph passes. + recorder = _Recorder() + runner = _runner(cuda_graphs_enabled=True) + runner.mapping = SimpleNamespace(cp_size=1, has_cp_helix=lambda: False) + runner.guided_decoder = None + # No context warmups: the attention, general, autotuner and Mamba phases + # are skipped, which leaves the capture sequence as the only forward. + runner._context_warmups_supported = False + stubbed_methods = ( + "_prewarm_engram_kernels", + "_prewarm_csa2_sparse_indexer", + "_prewarm_cute_dsl_indexer_q", + "_ensure_dsa_attn_metadata_for_warmup", + "_warmup_dg_paged_mqa_logits_metadata", + "_warmup_cute_dsl_radix_topk", + "_freeze_eager_workspace_floor", + ) + with contextlib.ExitStack() as stack: + stack.enter_context(_patched(recorder)) + for name in stubbed_methods: + stack.enter_context(mock.patch.object(DecoderRunner, name, autospec=True)) + stack.enter_context( + mock.patch.object( + DecoderRunner, + "_agree_warmup_flag", + autospec=True, + side_effect=lambda self, flag: flag, + ) + ) + for name in ("warmup_sampling_module", "log_mem_snapshot", "AutoTuner"): + stack.enter_context(mock.patch.object(runner_module, name)) + runner._warmup_scheduled(resource_manager=None, kv_cache_manager=SimpleNamespace()) + + assert recorder.events == ["drain_and_verify", *_CAPTURE_UNDER_BUDGET]