From d0138510d327f5bd5f36283b9a74787aff81b4d1 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:10:51 +0000 Subject: [PATCH 01/15] [None][perf] KVCacheManagerV2: opt-in stream-ordered events, block-key chain memo and lazy slot scrub Three independent changes to the KV cache manager v2 admission path. Each is off by default; with its variable unset the manager behaves as before. - TLLM_KV_CACHE_MANAGER_V2_STREAM_ORDERED_EVENTS=1: page resume and close skip event waits and merges that the stream order already implies, and wait once per distinct event instead of once per page. - TLLM_KV_CACHE_MANAGER_V2_KEY_CHAIN_MEMO=1: the radix tree memoises the block-key chain of a token sequence, so the reuse probe, the first-block probe and the cache creation of one request hash its prompt once. - TLLM_KV_CACHE_MANAGER_V2_LAZY_SLOT_SCRUB=1: a slot allocator scrubs the events of recycled slots once per batch allocation instead of polling them per slot. The nanobind module exposes the switches and the event counters for the unit tests. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../kv_cache_manager_v2/blockRadixTree.cpp | 187 +++++++++++ .../kv_cache_manager_v2/blockRadixTree.h | 54 ++++ .../kv_cache_manager_v2/storage/core.cpp | 38 ++- .../kv_cache_manager_v2/storage/core.h | 3 + .../kv_cache_manager_v2/utils/cudaEvent.cpp | 304 +++++++++++++++++- .../kv_cache_manager_v2/utils/cudaEvent.h | 46 ++- .../batch_manager/kvCacheManagerV2.cpp | 31 ++ .../test_key_chain_memo.py | 152 +++++++++ .../test_stream_ordered_events.py | 202 ++++++++++++ 9 files changed, 1006 insertions(+), 11 deletions(-) create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_key_chain_memo.py create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_stream_ordered_events.py diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp index d38b7ec2d7ef..27a00a3f8d35 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp @@ -27,6 +27,7 @@ #include "tensorrt_llm/common/assert.h" #include #include +#include #include #include #include @@ -696,9 +697,127 @@ BlockRadixTree::BlockRadixTree( : mLifeCycles(lifeCycles) , mTokensPerBlock(tokensPerBlock) , mEventSink(std::move(eventSink)) + , mKeyChainMemo(tokensPerBlock) { } +// --------------------------------------------------------------------------- +// BlockKeyChainMemo +// --------------------------------------------------------------------------- + +bool keyChainMemoEnabled() noexcept +{ + static bool const enabled = []() + { + char const* value = std::getenv("TLLM_KV_CACHE_MANAGER_V2_KEY_CHAIN_MEMO"); + return value != nullptr && value[0] == '1' && value[1] == '\0'; + }(); + return enabled; +} + +BlockKeyChainMemo::BlockKeyChainMemo(int tokensPerBlock, size_t capacity) + : mTokensPerBlock(static_cast(tokensPerBlock)) + , mCapacity(capacity) +{ + TLLM_CHECK(tokensPerBlock > 0 && capacity > 0); +} + +namespace +{ +//! Number of leading full blocks on which `a` and `b` (both at least `numBlocks` blocks long) agree. +size_t commonFullBlocks(uint32_t const* a, uint32_t const* b, size_t numBlocks, size_t tokensPerBlock) +{ + // Compare in large chunks first (memcmp is vectorised), then find the first differing block. + constexpr size_t kBlocksPerChunk = 64; + size_t block = 0; + while (block < numBlocks) + { + size_t const n = std::min(kBlocksPerChunk, numBlocks - block); + size_t const offset = block * tokensPerBlock; + if (std::memcmp(a + offset, b + offset, n * tokensPerBlock * sizeof(uint32_t)) == 0) + { + block += n; + continue; + } + for (size_t i = 0; i < n; ++i) + { + size_t const o = (block + i) * tokensPerBlock; + if (std::memcmp(a + o, b + o, tokensPerBlock * sizeof(uint32_t)) != 0) + { + return block + i; + } + } + TLLM_CHECK_DEBUG_WITH_INFO(false, "chunk differs but no block does"); + return block; + } + return numBlocks; +} +} // namespace + +BlockKeyChainMemo::Hit BlockKeyChainMemo::find( + BlockKey const& rootKey, TokenIdExt const* tokens, size_t numFullBlocks) const +{ + static_assert(sizeof(TokenIdExt) == sizeof(uint32_t)); + std::vector> entries; + { + std::lock_guard const lock(mMutex); + entries = mEntries; + } + auto const* bits = reinterpret_cast(tokens); + Hit best; + for (auto it = entries.rbegin(); it != entries.rend(); ++it) + { + Entry const& entry = **it; + if (entry.rootKey != rootKey) + { + continue; + } + size_t const limit = std::min(numFullBlocks, entry.keys.size()); + if (limit <= best.numBlocks) + { + continue; + } + size_t const n = commonFullBlocks(bits, entry.tokenBits.data(), limit, mTokensPerBlock); + if (n > best.numBlocks) + { + best = Hit{*it, n}; + } + } + return best; +} + +void BlockKeyChainMemo::remember( + BlockKey const& rootKey, TokenIdExt const* tokens, Hit const& hit, std::vector newKeys) +{ + if (newKeys.empty()) + { + return; + } + auto entry = std::make_shared(); + entry->rootKey = rootKey; + size_t const numBlocks = hit.numBlocks + newKeys.size(); + entry->keys.reserve(numBlocks); + if (hit.numBlocks > 0) + { + entry->keys.assign(hit.entry->keys.begin(), hit.entry->keys.begin() + static_cast(hit.numBlocks)); + } + entry->keys.insert(entry->keys.end(), newKeys.begin(), newKeys.end()); + auto const* bits = reinterpret_cast(tokens); + entry->tokenBits.assign(bits, bits + numBlocks * mTokensPerBlock); + + std::lock_guard const lock(mMutex); + // The new chain extends the hit entry when it covered all of that entry: replace it. + if (hit.entry && hit.numBlocks == hit.entry->keys.size()) + { + std::erase(mEntries, hit.entry); + } + if (mEntries.size() >= mCapacity) + { + mEntries.erase(mEntries.begin()); + } + mEntries.push_back(std::move(entry)); +} + BlockRadixTree::~BlockRadixTree() { // Detach blocks leaf-first in O(1) extra space. Dropping mRoots directly would instead @@ -790,6 +909,12 @@ std::vector BlockRadixTree::matchTokenPath( { // Read-only: probeReuse() runs this under a shared lock, so probes may overlap. A root still // awaiting erasure is childless, so it yields no match. + if (keyChainMemoEnabled() + && (knownNoDigest + || std::none_of(tokens.begin(), tokens.end(), [](TokenIdExt const& t) { return t.isDigest(); }))) + { + return matchTokenPathMemo(reuseScope, tokens, enablePartialMatch); + } std::vector results; // Lazily compute one key per iteration — no wasted hashing on early miss. @@ -837,6 +962,68 @@ std::vector BlockRadixTree::matchTokenPath( return results; } +std::vector BlockRadixTree::matchTokenPathMemo( + ReuseScope const& reuseScope, TokenSpan tokens, bool enablePartialMatch) const +{ + // The same walk as matchTokenPath; only where each block key comes from differs. Keys of the + // memoised prefix are read from the memo, the rest are hashed exactly as the generator does + // (Hasher over the previous key and the block's tokens) and recorded afterwards. + std::vector results; + BlockKey const rootKey = RootBlock::makeKey(reuseScope); + auto rootIt = mRoots.find(rootKey); + if (rootIt == mRoots.end()) + return results; + + size_t const tokensPerBlock = static_cast(mTokensPerBlock); + size_t const numTokens = static_cast(tokens.size()); + auto const hit = mKeyChainMemo.find(rootKey, tokens.begin(), numTokens / tokensPerBlock); + std::vector newKeys; + + std::unordered_map> const* currentNext = &rootIt->second->next; + HalfOpenRange missedRange; + bool missed = false; + BlockKey digest = rootKey; + for (size_t ordinal = 0, beg = 0; beg < numTokens; ++ordinal, beg += tokensPerBlock) + { + size_t const end = std::min(beg + tokensPerBlock, numTokens); + if (ordinal < hit.numBlocks) + { + digest = hit.entry->keys[ordinal]; + } + else + { + Hasher h; + h.update(digest); + h.update(tokens.begin() + beg, end - beg, /*knownNoDigest=*/true); + digest = h.digest(); + if (end - beg == tokensPerBlock) + { + newKeys.push_back(digest); + } + } + auto blockIt = currentNext->find(digest); + if (blockIt == currentNext->end()) + { + missedRange = {beg, end}; + missed = true; + break; + } + Block* block = blockIt->second.get(); + results.push_back({block, static_cast(end - beg)}); + currentNext = &block->next; + } + mKeyChainMemo.remember(rootKey, tokens.begin(), hit, std::move(newKeys)); + + if (missed && enablePartialMatch) + { + auto [best, bestMatch] + = findBestPartialMatchInNextNodes(*currentNext, tokens.begin() + missedRange.beg, missedRange.length()); + if (best) + results.push_back({best, bestMatch}); + } + return results; +} + std::vector BlockRadixTree::pruneMatch( std::vector matched, std::optional ssmLcId) const { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h index bdae500113db..88fe59f6052e 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h @@ -31,6 +31,7 @@ #include #include #include +#include #include #include #include @@ -347,6 +348,54 @@ struct Block : NodeBase, EnableSharedFromThis std::shared_ptr mLastTokenDigest; }; +// --------------------------------------------------------------------------- +// BlockKeyChainMemo — exact memo of recent block-key chains (opt-in: +// TLLM_KV_CACHE_MANAGER_V2_KEY_CHAIN_MEMO=1, read once per process). +// +// A block key depends only on the reuse scope and the tokens up to the end of that block, so the +// chain of a prompt is the same in every radix walk of that prompt. One admission walks it +// several times (shortest-prefill-first probe, first-new-block probe, cache creation), each time +// paying one SHA-256 per block. The memo keeps the full-block keys of the last few walked +// sequences together with a copy of their tokens; a later walk reuses the keys of its longest +// token-for-token equal prefix and hashes only the rest. Only digest-free (plain id) token +// sequences are memoised, so equal token bits mean equal content. Thread-safe: probes walk the +// tree concurrently under the manager's shared lock. +// --------------------------------------------------------------------------- +[[nodiscard]] bool keyChainMemoEnabled() noexcept; + +class BlockKeyChainMemo +{ +public: + struct Entry + { + BlockKey rootKey; + std::vector tokenBits; // exactly keys.size() full blocks + std::vector keys; + }; + + struct Hit + { + std::shared_ptr entry; + size_t numBlocks = 0; // leading full blocks whose keys entry->keys supplies + }; + + explicit BlockKeyChainMemo(int tokensPerBlock, size_t capacity = 32); + + //! The entry sharing the longest prefix of full blocks with `tokens` (digest-free) under + //! `rootKey`. numBlocks is 0 when nothing matches. + [[nodiscard]] Hit find(BlockKey const& rootKey, TokenIdExt const* tokens, size_t numFullBlocks) const; + + //! Record the chain of `tokens`: the hit's leading keys followed by `newKeys`, the keys of the + //! next full blocks. A no-op when there is nothing new. + void remember(BlockKey const& rootKey, TokenIdExt const* tokens, Hit const& hit, std::vector newKeys); + +private: + size_t mTokensPerBlock; + size_t mCapacity; + mutable std::mutex mMutex; + std::vector> mEntries; // least recently recorded first +}; + // --------------------------------------------------------------------------- // BlockRadixTree — the global cache index. // next: reuseScope digest → RootBlock. @@ -433,6 +482,9 @@ class BlockRadixTree // knownNoDigest: from external text_only knowledge, never a scan (see Hasher::update). std::vector matchTokenPath( ReuseScope const& reuseScope, TokenSpan tokens, bool knownNoDigest, bool enablePartialMatch) const; + // matchTokenPath for a digest-free sequence, with block keys taken from / recorded in mKeyChainMemo. + std::vector matchTokenPathMemo( + ReuseScope const& reuseScope, TokenSpan tokens, bool enablePartialMatch) const; // Shorten `matched` to the prefix that is actually reusable. Passing // std::nullopt for `ssmLcId` skips the recurrent-snapshot constraint and // yields the attention-only prefix (used for numReusableTokensBeforeHybridPruning). @@ -448,6 +500,8 @@ class BlockRadixTree std::unordered_map> mRoots; mutable std::vector mPendingRootErases; + // Pure function of (scope, tokens), so it never goes stale as the tree changes. + mutable BlockKeyChainMemo mKeyChainMemo; }; // --------------------------------------------------------------------------- diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp index 62d51bb6c550..c89e73eaa308 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.cpp @@ -22,6 +22,7 @@ #include #include #include +#include #include #include #include @@ -86,6 +87,24 @@ SlotCount SlotAllocator::numOccupiedSlots() const noexcept return slotCountValueFromSize(mOccupiedMask.numSetBits()); } +namespace +{ +//! TLLM_KV_CACHE_MANAGER_V2_LAZY_SLOT_SCRUB=1 (read once): scrub recycled-slot events once per +//! allocateMultiple() instead of once per slot, and not on release(). Scrubbing only decides which +//! free slot is handed out first (ready ones are preferred); every user of a slot still waits for +//! its readyEvent, so this changes neither correctness nor contents, only how many cuEventQuery +//! calls a migration of N pages issues (about 2 instead of 2N). +bool lazySlotScrubEnabled() noexcept +{ + static bool const enabled = []() + { + char const* value = std::getenv("TLLM_KV_CACHE_MANAGER_V2_LAZY_SLOT_SCRUB"); + return value != nullptr && value[0] == '1' && value[1] == '\0'; + }(); + return enabled; +} +} // namespace + Slot SlotAllocator::allocate() { if (numFreeSlots() == 0) @@ -93,7 +112,11 @@ Slot SlotAllocator::allocate() throw OutOfPagesError("SlotAllocator: no free slots"); } scrubEvents(); + return takeFreeSlot(); +} +Slot SlotAllocator::takeFreeSlot() +{ Slot slot; if (mNumReadyRecycledSlots > 0) { @@ -130,6 +153,16 @@ std::vector SlotAllocator::allocateMultiple(SlotCount numSlots) } std::vector result; result.reserve(slotCountToSizeT(numSlots)); + if (lazySlotScrubEnabled()) + { + // numFreeSlots() >= numSlots was checked above, and each take consumes exactly one. + scrubEvents(); + for (SlotCount slotCount{0}; slotCount < numSlots; ++slotCount) + { + result.push_back(takeFreeSlot()); + } + return result; + } for (SlotCount slotCount{0}; slotCount < numSlots; ++slotCount) { result.push_back(allocate()); @@ -160,7 +193,10 @@ void SlotAllocator::release(Slot slot) { mOverflowSlots.push_back(std::move(slot)); } - scrubEvents(); + if (!lazySlotScrubEnabled()) + { + scrubEvents(); + } TLLM_CHECK_DEBUG(check()); } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h index 94b15dc97520..5fb2fea890fc 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/core.h @@ -189,6 +189,9 @@ class SlotAllocator private: void scrubEvents(); + // allocate() after its scrub: the next ready recycled slot, else a never-used slot, else the + // oldest recycled slot (its readyEvent still pending). Precondition: numFreeSlots() > 0. + Slot takeFreeSlot(); [[nodiscard]] bool check() const noexcept; SlotCount mCapacity; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.cpp index 3ca1d10b5525..5188db108807 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.cpp @@ -18,9 +18,234 @@ #include "kv_cache_manager_v2/utils/cudaEvent.h" #include "kv_cache_manager_v2/exceptions.h" +#include +#include + namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { +// --------------------------------------------------------------------------- +// Stream-ordered events: record registry and wait selection +// --------------------------------------------------------------------------- + +namespace +{ + +bool readStreamOrderedEventsEnv() noexcept +{ + char const* value = std::getenv("TLLM_KV_CACHE_MANAGER_V2_STREAM_ORDERED_EVENTS"); + return value != nullptr && value[0] == '1' && value[1] == '\0'; +} + +//! Identity of a stream for the stream-order rules, or 0 when it is unknown or must not be used: +//! the NULL / legacy / per-thread default stream handles name different streams in different +//! contexts or threads. cuStreamGetId is unique for the lifetime of the context, so a destroyed +//! stream whose handle value is reused does not alias the new one. +uint64_t streamIdentity(CUstream stream) noexcept +{ + if (stream == nullptr || stream == CU_STREAM_LEGACY || stream == CU_STREAM_PER_THREAD) + { + return 0; + } + unsigned long long id = 0; + if (cuStreamGetId(stream, &id) != CUDA_SUCCESS) + { + return 0; + } + return static_cast(id) + 1; // 0 stays reserved for "unknown" +} + +struct EventRecordInfo +{ + uint64_t streamId = 0; // 0: unknown, never elided + uint64_t seq = 0; +}; + +//! The latest record of every pooled CUevent. Pooled events are only recorded through +//! CachedCudaEvent, and a CUevent is recorded again only after its PooledEvent has retired it, so +//! the entry for a handle loaded from a live CachedCudaEvent describes that event's record. +class EventRecordRegistry +{ +public: + static EventRecordRegistry& instance() + { + static EventRecordRegistry registry; + return registry; + } + + //! Record `event` on `stream`. The sequence number is taken under the same mutex as the + //! record, so of two records on one stream the larger number is the later one in stream + //! order, even when they come from different threads. + void record(CUevent event, CUstream stream) + { + uint64_t const streamId = streamIdentity(stream); + std::lock_guard const lock(mMutex); + cuCheck(cuEventRecord(event, stream)); + mRecords[event] = EventRecordInfo{streamId, ++mSeq}; + } + + void lookup(CUevent const* events, size_t count, EventRecordInfo* out) const + { + std::lock_guard const lock(mMutex); + for (size_t i = 0; i < count; ++i) + { + auto const it = mRecords.find(events[i]); + out[i] = it == mRecords.end() ? EventRecordInfo{} : it->second; + } + } + +private: + mutable std::mutex mMutex; + std::unordered_map mRecords; + uint64_t mSeq = 0; +}; + +struct AtomicStreamOrderedCounters +{ + std::atomic waitsIssued{0}; + std::atomic waitsElided{0}; + std::atomic mergesWithoutStream{0}; + std::atomic mergesWithStream{0}; +}; + +AtomicStreamOrderedCounters gStreamOrderedCounters; + +//! Indices into `infos` of the events that a stream with identity `waiterId` (0: none) still has +//! to wait for: every event of unknown stream, plus the latest event of each known stream other +//! than the waiter's own. Unknown-stream events come first, in input order, then one event per +//! stream in first-seen order. +void selectStreamOrderedWaits(uint64_t waiterId, std::vector const& infos, std::vector& out) +{ + struct Latest + { + uint64_t streamId; + uint64_t seq; + size_t index; + }; + + out.clear(); + thread_local std::vector latest; + thread_local std::unordered_map slotOfStream; + latest.clear(); + if (!slotOfStream.empty()) + { + // clear() resets every bucket, so skip it in the common case where the map was not used. + slotOfStream.clear(); + } + constexpr size_t kLinearSearchLimit = 16; + for (size_t i = 0; i < infos.size(); ++i) + { + auto const& info = infos[i]; + if (info.streamId == 0) + { + out.push_back(i); + continue; + } + if (info.streamId == waiterId) + { + continue; + } + size_t slot = latest.size(); + if (latest.size() <= kLinearSearchLimit) + { + for (size_t j = 0; j < latest.size(); ++j) + { + if (latest[j].streamId == info.streamId) + { + slot = j; + break; + } + } + } + else if (auto const it = slotOfStream.find(info.streamId); it != slotOfStream.end()) + { + slot = it->second; + } + if (slot == latest.size()) + { + latest.push_back({info.streamId, info.seq, i}); + if (latest.size() == kLinearSearchLimit + 1) + { + for (size_t j = 0; j < latest.size(); ++j) + { + slotOfStream.emplace(latest[j].streamId, j); + } + } + else if (latest.size() > kLinearSearchLimit + 1) + { + slotOfStream.emplace(info.streamId, slot); + } + } + else if (info.seq > latest[slot].seq) + { + latest[slot].seq = info.seq; + latest[slot].index = i; + } + } + for (auto const& entry : latest) + { + out.push_back(entry.index); + } +} + +} // namespace + +bool streamOrderedEventsEnabled() noexcept +{ + static bool const enabled = readStreamOrderedEventsEnv(); + return enabled; +} + +StreamOrderedEventCounters streamOrderedEventCounters() noexcept +{ + return StreamOrderedEventCounters{gStreamOrderedCounters.waitsIssued.load(std::memory_order_relaxed), + gStreamOrderedCounters.waitsElided.load(std::memory_order_relaxed), + gStreamOrderedCounters.mergesWithoutStream.load(std::memory_order_relaxed), + gStreamOrderedCounters.mergesWithStream.load(std::memory_order_relaxed)}; +} + +void streamWaitEvents(CudaStream stream, std::vector events) +{ + events.erase(std::remove(events.begin(), events.end(), nullptr), events.end()); + if (events.empty()) + { + return; + } + auto const cuStream = reinterpret_cast(stream); + if (!streamOrderedEventsEnabled()) + { + std::sort(events.begin(), events.end()); + events.erase(std::unique(events.begin(), events.end()), events.end()); + for (CUevent h : events) + { + cuCheck(cuStreamWaitEvent(cuStream, h, 0)); + } + return; + } + // Select first: copies of one event share its stream and sequence number, so the selection + // already collapses the duplicates of known streams and only the few remaining events need + // the sort-and-unique pass. + thread_local std::vector infos; + thread_local std::vector selected; + infos.resize(events.size()); + EventRecordRegistry::instance().lookup(events.data(), events.size(), infos.data()); + selectStreamOrderedWaits(streamIdentity(cuStream), infos, selected); + std::vector waits; + waits.reserve(selected.size()); + for (size_t const index : selected) + { + waits.push_back(events[index]); + } + std::sort(waits.begin(), waits.end()); + waits.erase(std::unique(waits.begin(), waits.end()), waits.end()); + gStreamOrderedCounters.waitsIssued.fetch_add(waits.size(), std::memory_order_relaxed); + gStreamOrderedCounters.waitsElided.fetch_add(events.size() - waits.size(), std::memory_order_relaxed); + for (CUevent h : waits) + { + cuCheck(cuStreamWaitEvent(cuStream, h, 0)); + } +} + // --------------------------------------------------------------------------- // CudaEventPool / CudaStreamPool singleton implementations // --------------------------------------------------------------------------- @@ -78,7 +303,14 @@ CachedCudaEvent::CachedCudaEvent(CudaStream stream) noexcept [&]() { mEvent = std::make_shared(CudaEventPool::instance().get()); - cuCheck(cuEventRecord(mEvent->load(), reinterpret_cast(stream))); + if (streamOrderedEventsEnabled()) + { + EventRecordRegistry::instance().record(mEvent->load(), reinterpret_cast(stream)); + } + else + { + cuCheck(cuEventRecord(mEvent->load(), reinterpret_cast(stream))); + } }); } @@ -120,7 +352,19 @@ void CachedCudaEvent::waitInStream(CudaStream stream) const { return; } - cuCheck(cuStreamWaitEvent(reinterpret_cast(stream), event, 0)); + auto const cuStream = reinterpret_cast(stream); + if (streamOrderedEventsEnabled()) + { + EventRecordInfo info; + EventRecordRegistry::instance().lookup(&event, 1, &info); + if (info.streamId != 0 && info.streamId == streamIdentity(cuStream)) + { + gStreamOrderedCounters.waitsElided.fetch_add(1, std::memory_order_relaxed); + return; + } + gStreamOrderedCounters.waitsIssued.fetch_add(1, std::memory_order_relaxed); + } + cuCheck(cuStreamWaitEvent(cuStream, event, 0)); } void CachedCudaEvent::close() const noexcept @@ -184,6 +428,62 @@ CachedCudaEvent mergeEvents(std::vector& events) return CachedCudaEvent::makeNull(); if (numLive == 1) return std::move(*onlyLive); + if (streamOrderedEventsEnabled()) + { + // Keep the latest event of every known stream (it completes after every other event + // recorded there) and each event of unknown stream, then drop those already complete. + // Only when more than one is left does the merge need a temporary stream. + std::vector handles; + std::vector liveIndex; + handles.reserve(numLive); + liveIndex.reserve(numLive); + for (size_t i = 0; i < events.size(); ++i) + { + CUevent const h = events[i].handle(); + if (h != nullptr) + { + handles.push_back(h); + liveIndex.push_back(i); + } + } + std::vector infos(handles.size()); + EventRecordRegistry::instance().lookup(handles.data(), handles.size(), infos.data()); + std::vector selected; + selectStreamOrderedWaits(/*waiterId=*/0, infos, selected); + std::vector keep; + keep.reserve(selected.size()); + for (size_t const s : selected) + { + // Copies of one unknown-stream event share a handle; keep it once. + bool const duplicate + = std::any_of(keep.begin(), keep.end(), [&](size_t k) { return handles[k] == handles[s]; }); + if (!duplicate) + { + keep.push_back(s); + } + } + if (keep.size() > 1) + { + std::erase_if(keep, [&](size_t k) { return events[liveIndex[k]].queryComplete(); }); + } + if (keep.size() <= 1) + { + gStreamOrderedCounters.mergesWithoutStream.fetch_add(1, std::memory_order_relaxed); + return keep.empty() ? CachedCudaEvent::makeNull() : std::move(events[liveIndex[keep.front()]]); + } + gStreamOrderedCounters.mergesWithStream.fetch_add(1, std::memory_order_relaxed); + std::vector priors; + priors.reserve(keep.size()); + for (size_t const k : keep) + { + priors.push_back(handles[k]); + } + TemporaryCudaStream tempStream(std::move(priors)); + { + auto scope = tempStream.enter(); + } + return tempStream.takeFinishEvent(); + } // Multiple live events: merge via TemporaryCudaStream. Closed events yield a null handle, // which streamWaitEvents drops. std::vector priors; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.h index 2c60f5e6a2ba..614f14e5b159 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/cudaEvent.h @@ -23,6 +23,7 @@ #include #include +#include #include #include #include @@ -378,15 +379,44 @@ class CachedCudaEvent // Wait for all events on the given stream, skipping nulls and issuing one wait per distinct // event. Mirrors Python's stream_wait_events(), which converts to set() before iterating. // Waiting does not consume an event, so raw CUevent values suffice and a caller-owned -// CachedCudaEvent need not outlive the call. -inline void streamWaitEvents(CudaStream stream, std::vector events) +// CachedCudaEvent need not outlive the call. With stream-ordered events enabled (below), waits +// that stream order already implies are not issued. +void streamWaitEvents(CudaStream stream, std::vector events); + +// --------------------------------------------------------------------------- +// Stream-ordered events (opt-in: TLLM_KV_CACHE_MANAGER_V2_STREAM_ORDERED_EVENTS=1, read once). +// +// Every CachedCudaEvent record is registered with the identity of the stream it was recorded on +// (cuStreamGetId) and a process-wide sequence number, both taken under the mutex that also +// covers the cuEventRecord, so two registered records on one stream compare in stream order. +// CUDA streams execute in order, which gives two rules: +// - a stream never needs to wait for an event recorded earlier on that same stream; +// - of several events recorded on one stream, the latest one completes last. +// streamWaitEvents(), CachedCudaEvent::waitInStream() and mergeEvents() apply them: they skip +// waits that the first rule makes redundant, keep only the latest event per stream, and +// mergeEvents() also drops events that have already completed before it falls back to a +// temporary merge stream. The GPU orderings that result are the same as without the rules. +// Records on the NULL / legacy / per-thread default streams, or on a stream whose identity +// cannot be read, are never elided. +// +// Typical effect: a request's pages are written and released on one execution stream, so when +// a cache is closed every page's finish events reduce to that stream's latest event (no +// temporary stream per page), and a later resume on the same stream issues no per-page waits. +// --------------------------------------------------------------------------- +[[nodiscard]] bool streamOrderedEventsEnabled() noexcept; + +//! Counters of the stream-ordered event paths (only counted while enabled), for tests and +//! diagnostics. A wait helper adds every non-null handle it is given (duplicates included) to +//! either waitsIssued or waitsElided. +struct StreamOrderedEventCounters { - events.erase(std::remove(events.begin(), events.end(), nullptr), events.end()); - std::sort(events.begin(), events.end()); - events.erase(std::unique(events.begin(), events.end()), events.end()); - for (CUevent h : events) - cuCheck(cuStreamWaitEvent(reinterpret_cast(stream), h, 0)); -} + uint64_t waitsIssued = 0; + uint64_t waitsElided = 0; + uint64_t mergesWithoutStream = 0; + uint64_t mergesWithStream = 0; +}; + +[[nodiscard]] StreamOrderedEventCounters streamOrderedEventCounters() noexcept; // --------------------------------------------------------------------------- // CachedCudaStream — pooled non-blocking CUstream. diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 047f6a8452a0..7171d179bba3 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -2131,6 +2131,37 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def("close", &kv::CachedCudaEvent::close) .def("is_closed", &kv::CachedCudaEvent::isClosed); cachedCudaEvent.attr("NULL") = kv::CachedCudaEvent::makeNull(); + mIntrospection.def("stream_ordered_events_enabled", &kv::streamOrderedEventsEnabled, + "Whether TLLM_KV_CACHE_MANAGER_V2_STREAM_ORDERED_EVENTS=1 was set when the process first used it."); + mIntrospection.def( + "stream_ordered_event_counters", + []() + { + auto const c = kv::streamOrderedEventCounters(); + nb::dict d; + d["waits_issued"] = nb::cast(c.waitsIssued); + d["waits_elided"] = nb::cast(c.waitsElided); + d["merges_without_stream"] = nb::cast(c.mergesWithoutStream); + d["merges_with_stream"] = nb::cast(c.mergesWithStream); + return d; + }, + "Process-wide counters of the stream-ordered event paths (counted only while enabled)."); + mIntrospection.def( + "stream_wait_events", + [](kv::CudaStream stream, std::vector const& events) + { + std::vector handles; + handles.reserve(events.size()); + for (auto const& ev : events) + { + handles.push_back(ev.handle()); + } + kv::streamWaitEvents(stream, std::move(handles)); + }, + nb::arg("stream"), nb::arg("events"), nb::call_guard()); + mIntrospection.def( + "merge_events", [](std::vector events) { return kv::mergeEvents(events); }, + nb::arg("events"), nb::call_guard()); mIntrospection.def( "make_test_block", [](kv::KvCacheManager& manager, nb::object tokenObject, std::vector coveragePerLc, nb::object parentObject, diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_key_chain_memo.py b/tests/unittest/kv_cache_manager_v2_tests/test_key_chain_memo.py new file mode 100644 index 000000000000..bbe1442be7c5 --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_key_chain_memo.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Block-key chain memo (TLLM_KV_CACHE_MANAGER_V2_KEY_CHAIN_MEMO) must not change any match. + +The switch is read once per process, so one deterministic scenario runs in two child processes, +with the memo off and on, and the transcripts of every probe and cache creation must be equal. +The scenario commits conversations that share a system prefix, then queries prefixes, extensions, +mid-block divergences, other reuse scopes, partial last blocks and int32 inputs while the tree +keeps changing (new commits) between queries, so memo entries outlive the tree states they were +recorded in. +""" + +import json +import os +import subprocess +import sys + +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +_ENV = "TLLM_KV_CACHE_MANAGER_V2_KEY_CHAIN_MEMO" + + +def _scenario() -> list: + sys.path.insert(0, _TEST_DIR) + import random + + import numpy as np + from cuda_test_utils import TemporaryCudaStream, init_cuda_once + from test_kv_cache_manager_v2 import create_config + + from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager, ReuseScope, TokenId + + init_cuda_once() + tpb = 4 + manager = KVCacheManager( + create_config( + tokens_per_block=tpb, + gpu_quota=64 << 20, + host_quota=0, + disk_quota=0, + num_layers=2, + window_size=None, + sink_tokens=0, + ) + ) + rng = random.Random(7) + transcript = [] + + def tok(seq): + return [TokenId(t) for t in seq] + + def commit(seq, scope): + with TemporaryCudaStream([]) as holder: + kv = manager.create_kv_cache(scope, tok(seq)) + assert kv.resume(holder.handle) + assert kv.resize(len(seq)) + done = kv.num_committed_tokens + if done < len(seq): + kv.commit(tok(seq[done:])) + kv.close() + holder.take_finish_event().synchronize() + + def query(seq, scope, as_array): + data = np.asarray(seq, dtype=np.int32) if as_array else tok(seq) + reuse = manager.probe_reuse(scope, data) + key = manager.probe_first_new_block_key(scope, data) + kv = manager.create_kv_cache(scope, data) + record = [ + len(seq), + reuse, + None if key is None else bytes(key).hex(), + kv.num_committed_tokens, + list(kv.cached_tokens_by_level), + ] + kv.close() + transcript.append(record) + + scopes = [ReuseScope(), ReuseScope(salt=11), ReuseScope(lora_id=3)] + system = [rng.randrange(1, 5000) for _ in range(37)] + convs = [] + for i in range(6): + conv = system + [rng.randrange(1, 5000) for _ in range(rng.randrange(5, 60))] + convs.append(conv) + commit(conv, scopes[i % 2]) + for step in range(120): + conv = convs[rng.randrange(len(convs))] + scope = scopes[rng.randrange(len(scopes))] + kind = rng.randrange(6) + if kind == 0: + seq = conv[: rng.randrange(1, len(conv) + 1)] + elif kind == 1: + seq = conv + [rng.randrange(1, 5000) for _ in range(rng.randrange(1, 30))] + elif kind == 2: + cut = rng.randrange(1, len(conv)) + seq = conv[:cut] + [5001 + rng.randrange(100)] + conv[cut + 1 :] + elif kind == 3: + seq = system[: rng.randrange(1, len(system) + 1)] + elif kind == 4: + seq = [rng.randrange(1, 5000) for _ in range(rng.randrange(1, 50))] + else: + seq = conv + query(seq, scope, as_array=bool(step % 2)) + if step % 9 == 0: + new_conv = conv + [rng.randrange(1, 5000) for _ in range(rng.randrange(3, 40))] + convs.append(new_conv) + commit(new_conv, scope) + manager.shutdown() + return transcript + + +def _run(enabled: bool) -> list: + env = dict(os.environ) + if enabled: + env[_ENV] = "1" + else: + env.pop(_ENV, None) + proc = subprocess.run( + [sys.executable, os.path.abspath(__file__), "--worker"], + env=env, + capture_output=True, + text=True, + timeout=600, + check=False, + ) + assert proc.returncode == 0, proc.stdout[-4000:] + proc.stderr[-4000:] + line = [ln for ln in proc.stdout.splitlines() if ln.startswith("RESULT ")][-1] + return json.loads(line[len("RESULT ") :]) + + +def test_memo_does_not_change_matches(): + off = _run(False) + on = _run(True) + assert len(off) == len(on) == 120 + # The scenario must exercise real reuse, not only misses. + assert sum(1 for r in off if r[1] > 0) > 40 + for i, (a, b) in enumerate(zip(off, on)): + assert a == b, (i, a, b) + + +if __name__ == "__main__" and "--worker" in sys.argv: + print("RESULT " + json.dumps(_scenario()), flush=True) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_stream_ordered_events.py b/tests/unittest/kv_cache_manager_v2_tests/test_stream_ordered_events.py new file mode 100644 index 000000000000..4264427fd55e --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_stream_ordered_events.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Stream-ordered KV-cache events (TLLM_KV_CACHE_MANAGER_V2_STREAM_ORDERED_EVENTS). + +The switch is read once per process, so every case runs in a child process with the variable set +or unset. Each case holds a producer stream open with a host gate and then checks two things: +the ordering the wait / merge helpers establish (a consumer marker must not complete while the +gate is held, and must complete after it is released), and, with the switch on, which waits and +merges the helpers skipped (the process-wide counters). +""" + +import json +import os +import subprocess +import sys + +import pytest + +_TEST_DIR = os.path.dirname(os.path.abspath(__file__)) +_ENV = "TLLM_KV_CACHE_MANAGER_V2_STREAM_ORDERED_EVENTS" + + +def _introspection(): + try: + from bindings.internal.batch_manager.kv_cache_manager_v2 import _introspection + except ImportError: + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2 import _introspection + return _introspection + + +def _counters_delta(intro, before): + after = intro.stream_ordered_event_counters() + return {k: after[k] - before[k] for k in after} + + +def _worker() -> dict: + """Run every case in this process and return their observations.""" + sys.path.insert(0, _TEST_DIR) + from cuda_test_utils import CachedCudaStream, init_cuda_once + from kernels import HostGate + + init_cuda_once() + intro = _introspection() + event = intro.CachedCudaEvent + out = {"enabled": intro.stream_ordered_events_enabled()} + a_holder, b_holder, c_holder = CachedCudaStream(), CachedCudaStream(), CachedCudaStream() + a, b, c = int(a_holder.handle), int(b_holder.handle), int(c_holder.handle) + + def sync_all(): + for h in (a_holder, b_holder, c_holder): + h.synchronize() + + # 1. A stream waiting for its own earlier event. + before = intro.stream_ordered_event_counters() + e = event(a) + intro.stream_wait_events(a, [e, e]) + e.wait_in_stream(a) + sync_all() + out["same_stream"] = _counters_delta(intro, before) + + # 2. Two events of one gated stream: the consumer must follow the later one. + gate = HostGate() + try: + early = event(a) + gate.block_stream(a) + late = event(a) + before = intro.stream_ordered_event_counters() + intro.stream_wait_events(b, [early, late]) + marker = event(b) + out["latest_blocked"] = not marker.query_complete() + gate.release() + marker.synchronize() + out["latest_done_after_release"] = marker.is_closed() + out["latest"] = _counters_delta(intro, before) + finally: + gate.close() + sync_all() + + # 3. Merge of two events of one gated stream. + gate = HostGate() + try: + early = event(a) + gate.block_stream(a) + late = event(a) + before = intro.stream_ordered_event_counters() + merged = intro.merge_events([early, late]) + merged.wait_in_stream(b) + marker = event(b) + out["merge_same_blocked"] = not marker.query_complete() + gate.release() + marker.synchronize() + out["merge_same"] = _counters_delta(intro, before) + finally: + gate.close() + sync_all() + + # 4. Merge of pending events of two gated streams: needs both. + gate_a, gate_b = HostGate(), HostGate() + try: + gate_a.block_stream(a) + gate_b.block_stream(b) + ea, eb = event(a), event(b) + before = intro.stream_ordered_event_counters() + merged = intro.merge_events([ea, eb]) + merged.wait_in_stream(c) + marker = event(c) + gate_a.release() + a_holder.synchronize() + out["merge_cross_blocked_on_b"] = not marker.query_complete() + gate_b.release() + marker.synchronize() + out["merge_cross"] = _counters_delta(intro, before) + finally: + gate_a.close() + gate_b.close() + sync_all() + + # 5. Merge of a pending event with a completed (not yet closed) event of another stream. + gate = HostGate() + try: + eb = event(b) + b_holder.synchronize() + gate.block_stream(a) + ea = event(a) + before = intro.stream_ordered_event_counters() + merged = intro.merge_events([eb, ea]) + merged.wait_in_stream(c) + marker = event(c) + out["merge_done_blocked"] = not marker.query_complete() + gate.release() + marker.synchronize() + out["merge_done"] = _counters_delta(intro, before) + finally: + gate.close() + sync_all() + return out + + +def _run(enabled: bool) -> dict: + env = dict(os.environ) + if enabled: + env[_ENV] = "1" + else: + env.pop(_ENV, None) + proc = subprocess.run( + [sys.executable, os.path.abspath(__file__), "--worker"], + env=env, + capture_output=True, + text=True, + timeout=600, + check=False, + ) + assert proc.returncode == 0, proc.stdout[-4000:] + proc.stderr[-4000:] + line = [ln for ln in proc.stdout.splitlines() if ln.startswith("RESULT ")][-1] + return json.loads(line[len("RESULT ") :]) + + +@pytest.mark.parametrize("enabled", [False, True], ids=["off", "on"]) +def test_orderings_hold(enabled): + r = _run(enabled) + assert r["enabled"] is enabled + assert r["latest_blocked"] and r["latest_done_after_release"] + assert r["merge_same_blocked"] + assert r["merge_cross_blocked_on_b"] + assert r["merge_done_blocked"] + if not enabled: + for case in ("same_stream", "latest", "merge_same", "merge_cross", "merge_done"): + assert all(v == 0 for v in r[case].values()), (case, r[case]) + + +def test_redundant_waits_and_merges_are_skipped(): + r = _run(True) + # Own-stream waits: two copies through the list helper, one through wait_in_stream. + assert r["same_stream"]["waits_issued"] == 0 + assert r["same_stream"]["waits_elided"] == 3 + # Two events of one other stream: only the later one is waited for. + assert r["latest"]["waits_issued"] == 1 + assert r["latest"]["waits_elided"] == 1 + # One stream: the later event is the merge, no temporary stream. + assert r["merge_same"]["merges_without_stream"] == 1 + assert r["merge_same"]["merges_with_stream"] == 0 + # Two pending streams: a real merge. + assert r["merge_cross"]["merges_with_stream"] == 1 + # A completed event is dropped, so the pending one is the merge. + assert r["merge_done"]["merges_without_stream"] == 1 + assert r["merge_done"]["merges_with_stream"] == 0 + + +if __name__ == "__main__" and "--worker" in sys.argv: + print("RESULT " + json.dumps(_worker()), flush=True) From 41442a5ce12f683498306f9d16655531ede770bd Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:10:53 +0000 Subject: [PATCH 02/15] [None][feat] NVFP4 cold-page KV: optional FP8 global scale and uncompressed layer types; DFlash without a separate draft KV cache ColdPageQuantizationCompressionConfig gains two fields, both defaulting to the previous behaviour: - global_scale: identity (default) or fp8_block_amax. For an FP8 KV cache without a ModelOpt scale, fp8_block_amax uses a global scale of 6, which maps each block's E4M3 absolute maximum onto an exact E4M3 block scale. - uncompressed_layer_types: layer types (sliding_window, full_attention) whose pages stay uncompressed in the cold tier. The KV cache compression check now also admits DFlash when it has no separate draft KV cache: its drafter keeps the context K/V in worker-owned buffers, so the codec only sees target pages. The llm_args golden manifest records the two new fields. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../nvfp4_quantization.py | 62 ++++- tensorrt_llm/_torch/pyexecutor/_util.py | 11 +- tensorrt_llm/llmapi/llm_args.py | 26 ++ .../usage/llm_args_golden_manifest.json | 18 ++ ...st_nvfp4_cold_page_fp8_global_scale_gpu.py | 165 ++++++++++++ .../test_quantization_for_cold_page.py | 237 +++++++++++++++++- tests/unittest/llmapi/test_llm_args.py | 2 + 7 files changed, 508 insertions(+), 13 deletions(-) create mode 100644 tests/unittest/_torch/kv_cache_compression/test_nvfp4_cold_page_fp8_global_scale_gpu.py diff --git a/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/nvfp4_quantization.py b/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/nvfp4_quantization.py index 4b7638d727e0..e9276620ef7b 100644 --- a/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/nvfp4_quantization.py +++ b/tensorrt_llm/_torch/kv_cache_compression/quantization_for_cold_page/nvfp4_quantization.py @@ -31,6 +31,9 @@ _LayerScales = dict[str, _ScalePair] _IDENTITY_NVFP4_SCALE: _ScalePair = (1.0, 1.0) +# E4M3 values are at most 448 = 6 * 448 / 6, so a global scale of 6 maps each +# block's E4M3 absolute maximum onto an exact E4M3 block scale and E2M1 value 6. +_FP8_BLOCK_AMAX_NVFP4_SCALE: _ScalePair = (6.0, 1.0 / 6.0) _MODEL_OPT_LANGUAGE_KV_SCALE_KEY = re.compile( r"^model(?:\.language_model)?\.layers\.(?P\d+)\.self_attn\." r"(?P[kv])_proj\.(?P=kind)_scale$" @@ -212,10 +215,16 @@ def __init__( self._model_scales = _load_modelopt_nvfp4_scales(config.scale_checkpoint_path) self._skip_rope_quantization = bool(config.skip_rope_quantization) self._nvfp4_residual_dim = config.nvfp4_residual_dim + self._global_scale = config.global_scale + self._uncompressed_layer_types = frozenset(config.uncompressed_layer_types) model_type = getattr(pretrained_config, "model_type", None) architectures = getattr(pretrained_config, "architectures", None) or [] if architectures and architectures[0] == "GlmMoeDsaForCausalLM": model_type = "glm_moe_dsa" + if self._uncompressed_layer_types and model_type == "deepseek_v4": + raise NotImplementedError( + "uncompressed_layer_types is not supported for DeepSeek-V4 cold pages" + ) if self._skip_rope_quantization and model_type not in _SKIP_ROPE_QUANTIZATION_MODEL_TYPES: logger.warning( "skip_rope_quantization is supported for model types " @@ -224,6 +233,49 @@ def __init__( ) self._skip_rope_quantization = False + def _default_scale(self, runtime_dtype: DataType, *, is_draft: bool) -> _ScalePair: + """Global scale for buffers that have no ModelOpt scale.""" + + if self._global_scale == "identity": + return _IDENTITY_NVFP4_SCALE + if runtime_dtype == DataType.FP8: + return _FP8_BLOCK_AMAX_NVFP4_SCALE + if is_draft: + logger.warning( + f"global_scale={self._global_scale!r} needs an FP8 KV cache; the " + f"{runtime_dtype} draft KV cache keeps identity global scales." + ) + return _IDENTITY_NVFP4_SCALE + raise ValueError( + f"global_scale={self._global_scale!r} requires an FP8 KV cache, not {runtime_dtype}" + ) + + def _uncompressed_layer_ids( + self, attention_layers: Sequence["AttentionLayerConfig"] + ) -> frozenset[int]: + """Attention layers left out of the codec, so their lifecycles use the lossless copy.""" + + if not self._uncompressed_layer_types or not attention_layers: + return frozenset() + windows = { + int(layer.layer_id): ( + math.inf if layer.sliding_window_size is None else layer.sliding_window_size + ) + for layer in attention_layers + } + longest = max(windows.values()) + result = frozenset( + layer_id + for layer_id, window in windows.items() + if ("sliding_window" if window < longest else "full_attention") + in self._uncompressed_layer_types + ) + logger.info( + f"NVFP4 cold pages keep {len(result)} of {len(windows)} Attention layers " + f"({sorted(self._uncompressed_layer_types)}) in the active-cache dtype" + ) + return result + def _calculate_quantized_range( self, row_elements: int, @@ -320,6 +372,7 @@ def _build_deepseek_v4_layer_layouts( runtime_type: int, is_draft: bool, skip_rope: bool, + default_scale: _ScalePair = _IDENTITY_NVFP4_SCALE, ) -> dict[int, _Nvfp4LayerLayout | None]: """Build layouts for DeepSeek-V4 lifecycles; None selects lossless fallback.""" @@ -409,7 +462,7 @@ def _build_deepseek_v4_layer_layouts( f"{configured_bytes} bytes; expected {raw_bytes} for an ordinary runtime Page" ) - scale = _IDENTITY_NVFP4_SCALE + scale = default_scale if not is_draft: model_scales = self._model_scales.get(model_layers[layer_id]) if model_scales: @@ -475,6 +528,8 @@ def build_codec_state( "NVFP4 cold-page compression supports FP16, BF16, or FP8 " f"Attention KV, not {runtime_dtype}" ) + default_scale = self._default_scale(runtime_dtype, is_draft=is_draft) + uncompressed_layers = self._uncompressed_layer_ids(attention_layers) # The codec holds only the target model's config, so it cannot locate RoPE # in a draft model's K vectors: draft KVCMs always quantize whole vectors. @@ -494,11 +549,14 @@ def build_codec_state( runtime_type=runtime_type if runtime_type is not None else 0, is_draft=is_draft, skip_rope=skip_rope, + default_scale=default_scale, ) layer_layouts = [] for layer in attention_layers: layer_id = int(layer.layer_id) + if layer_id in uncompressed_layers: + continue if layer_id in deepseek_v4_layouts: layout = deepseek_v4_layouts[layer_id] if layout is not None: @@ -520,7 +578,7 @@ def build_codec_state( if len(compressed_roles) == 2 and model_scales: scales = tuple(model_scales[role] for role in ("k", "v")) else: - scales = (_IDENTITY_NVFP4_SCALE,) * len(compressed_roles) + scales = (default_scale,) * len(compressed_roles) num_kv_heads = int(num_kv_heads_per_layer[layer_id]) tokens_per_page = int(cache_config.tokens_per_block) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 68e370b61ed7..84bfecc85ac9 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -4107,8 +4107,15 @@ def validate_kv_cache_compression_compatibility( f"{guidance}") mode = spec_config.spec_dec_mode if config.algorithm == "quantization_for_cold_page": - supported = mode.is_mtp_eagle_one_model() or mode.is_eagle3_one_model() - guidance = "one-model MTP-EAGLE or EAGLE3" + # A DFlash drafter without a separate draft KVCM keeps its context K/V + # in worker-owned buffers, so the codec only ever sees target pages. + dflash_without_draft_kvcm = ( + mode.is_dflash() + and not should_use_separate_draft_kv_cache(spec_config)) + supported = (mode.is_mtp_eagle_one_model() + or mode.is_eagle3_one_model() or dflash_without_draft_kvcm) + guidance = ("one-model MTP-EAGLE or EAGLE3, or DFlash without a " + "separate draft KV cache") else: supported = mode.is_mtp_one_model() or mode.is_eagle3_one_model() guidance = "one-model MTP or EAGLE3" diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 5a6c913bdfe1..4a00f249fd4f 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4282,6 +4282,32 @@ class ColdPageQuantizationCompressionConfig(KvCacheCompressionConfig): description= "Optional local ModelOpt NVFP4 checkpoint directory supplying per-layer " "K/V global scales. Omit it to use identity global scales.") + global_scale: Literal["identity", "fp8_block_amax"] = Field( + default="identity", + status="prototype", + description="NVFP4 global scale for cold pages without ModelOpt scales. " + "'identity' (default): global scale 1.0. 'fp8_block_amax': for an FP8 " + "active cache, a global scale of 6, so that each 16-value block scale is " + "the block's E4M3 absolute maximum itself. It needs no calibration and " + "cannot overflow. Requires an FP8 target KV cache; cannot be combined " + "with scale_checkpoint_path.") + uncompressed_layer_types: List[Literal[ + "sliding_window", "full_attention"]] = Field( + default_factory=list, + status="prototype", + description= + "Attention layer types whose cold pages keep the active-cache " + "dtype (copied without quantization). 'sliding_window': layers " + "whose attention window is shorter than the longest window in the " + "KV cache; 'full_attention': the other Attention layers. Empty " + "(default): every Attention layer is quantized.") + + @model_validator(mode='after') + def _validate_global_scale_source(self): + if self.global_scale != "identity" and self.scale_checkpoint_path: + raise ValueError( + "global_scale and scale_checkpoint_path are mutually exclusive") + return self def supports_block_reuse(self) -> bool: """Block reuse is unchanged because token identity is preserved.""" diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 69df5d605c6b..5e56fafa1b3d 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -559,6 +559,15 @@ "kind": "categorical", "path": "kv_cache_compression_config.eviction_mode" }, + { + "allowed_values": [ + "identity", + "fp8_block_amax" + ], + "capture_policy": "literal", + "kind": "categorical", + "path": "kv_cache_compression_config.global_scale" + }, { "capture_policy": "bool", "kind": "value", @@ -586,6 +595,15 @@ "kind": "value", "path": "kv_cache_compression_config.skip_rope_quantization" }, + { + "allowed_values": [ + "sliding_window", + "full_attention" + ], + "capture_policy": "list[literal]", + "kind": "categorical", + "path": "kv_cache_compression_config.uncompressed_layer_types" + }, { "capture_policy": "int", "kind": "value", diff --git a/tests/unittest/_torch/kv_cache_compression/test_nvfp4_cold_page_fp8_global_scale_gpu.py b/tests/unittest/_torch/kv_cache_compression/test_nvfp4_cold_page_fp8_global_scale_gpu.py new file mode 100644 index 000000000000..b68d9c8719fb --- /dev/null +++ b/tests/unittest/_torch/kv_cache_compression/test_nvfp4_cold_page_fp8_global_scale_gpu.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +"""Single-GPU round trips of FP8 Pages through NVFP4 cold Pages for each global-scale policy.""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.kv_cache_compression.quantization_for_cold_page.nvfp4_quantization import ( + Nvfp4ColdPageQuantizationCompression, +) +from tensorrt_llm._torch.pyexecutor.resource_manager import DataType +from tensorrt_llm._utils import is_sm_100f +from tensorrt_llm.llmapi.llm_args import ColdPageQuantizationCompressionConfig +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionLayerConfig, BufferConfig + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not is_sm_100f(), + reason="NVFP4 cold-page kernels require an SM100-family GPU", +) + +_HEADS, _TOKENS, _HEAD_DIM, _PAGES = 4, 32, 64, 4 + + +def _fp8_pages(block_amax_log2: tuple[float, float], seed: int = 0) -> torch.Tensor: + """FP8 Pages [pages, heads, tokens, head_dim] whose 16-value blocks span a range of magnitudes.""" + + generator = torch.Generator(device="cuda").manual_seed(seed) + shape = (_PAGES, _HEADS, _TOKENS, _HEAD_DIM // 16, 16) + low, high = block_amax_log2 + block_amax = torch.exp2( + torch.empty(shape[:-1] + (1,), device="cuda").uniform_(low, high, generator=generator) + ) + values = torch.empty(shape, device="cuda").uniform_(-1.0, 1.0, generator=generator) + values = values / values.abs().amax(-1, keepdim=True) * block_amax + return values.reshape(_PAGES, _HEADS, _TOKENS, _HEAD_DIM).to(torch.float8_e4m3fn) + + +class _Codec: + """One FP8 K/V layer, encoded to and decoded from device cold scratch.""" + + def __init__(self, global_scale: str, key: torch.Tensor, value: torch.Tensor) -> None: + self.buffers = {"key": key, "value": value} + page_bytes = key[0].numel() + cache_config = SimpleNamespace( + tokens_per_block=_TOKENS, + layers=( + AttentionLayerConfig( + layer_id=0, + buffers=[BufferConfig(role=role, size=page_bytes) for role in self.buffers], + ), + ), + ) + self.provider = Nvfp4ColdPageQuantizationCompression( + ColdPageQuantizationCompressionConfig(global_scale=global_scale), + pretrained_config=SimpleNamespace(model_type="muse_spark"), + ) + self.state = self.provider.build_codec_state( + cache_config, + runtime_dtype=DataType.FP8, + pp_layers=(0,), + num_kv_heads_per_layer=(_HEADS,), + head_dim_per_layer=(_HEAD_DIM,), + ) + lifecycle = SimpleNamespace( + layers={ + 0: { + role: SimpleNamespace( + raw_base=buffer.data_ptr(), raw_slot_bytes=page_bytes, raw_bytes=page_bytes + ) + for role, buffer in self.buffers.items() + } + } + ) + self.provider.configure(self.state, [lifecycle]) + self.cold_page_bytes = self.state.lifecycle_metadata[0].cold_page_bytes + self.cold = torch.zeros(_PAGES * self.cold_page_bytes, dtype=torch.uint8, device="cuda") + self.stream = torch.cuda.current_stream() + + def encode(self, gpu_to_cold: list[tuple[int, int]]) -> None: + pairs = torch.tensor([[cold, gpu] for gpu, cold in gpu_to_cold], dtype=torch.int32) + self.provider.encode_cold_pages( + self.state, + 0, + self.cold.data_ptr(), + pairs.data_ptr(), + len(pairs), + self.stream.cuda_stream, + ) + self.stream.synchronize() + + def decode(self, cold_to_gpu: list[tuple[int, int]]) -> None: + pairs = torch.tensor([[gpu, cold] for cold, gpu in cold_to_gpu], dtype=torch.int32) + self.provider.decode_cold_pages( + self.state, + 0, + self.cold.data_ptr(), + pairs.data_ptr(), + len(pairs), + self.stream.cuda_stream, + ) + self.stream.synchronize() + + def cold_page(self, index: int) -> torch.Tensor: + return self.cold[index * self.cold_page_bytes : (index + 1) * self.cold_page_bytes].clone() + + +def _round_trip( + global_scale: str, source: dict[str, torch.Tensor] +) -> tuple[dict[str, torch.Tensor], _Codec]: + """Pages 0-1 hold the source; they are encoded to cold slots 0-1 and decoded into Pages 2-3.""" + + buffers = {role: torch.zeros_like(pages) for role, pages in source.items()} + for role, pages in source.items(): + buffers[role][:2] = pages[:2] + codec = _Codec(global_scale, buffers["key"], buffers["value"]) + codec.encode([(0, 0), (1, 1)]) + codec.decode([(0, 2), (1, 3)]) + return {role: buffer[2:].clone() for role, buffer in buffers.items()}, codec + + +def _relative_error(restored: torch.Tensor, source: torch.Tensor) -> float: + return ((restored.float() - source.float()).norm() / source.float().norm()).item() + + +def test_fp8_block_amax_restores_each_block_maximum_exactly() -> None: + source = {"key": _fp8_pages((-9.0, 8.8), seed=1), "value": _fp8_pages((-9.0, 8.8), seed=2)} + restored, _ = _round_trip("fp8_block_amax", source) + for role, pages in source.items(): + original_amax = pages[:2].float().reshape(-1, 16).abs().amax(-1) + restored_amax = restored[role].float().reshape(-1, 16).abs().amax(-1) + assert torch.equal(restored_amax, original_amax), role + + +def test_fp8_block_amax_is_more_accurate_than_identity_and_keeps_small_blocks() -> None: + source = {"key": _fp8_pages((-9.0, 8.8), seed=3), "value": _fp8_pages((-9.0, 8.8), seed=4)} + identity, _ = _round_trip("identity", source) + tuned, _ = _round_trip("fp8_block_amax", source) + for role, pages in source.items(): + assert _relative_error(tuned[role], pages[:2]) < _relative_error(identity[role], pages[:2]) + + # Identity block scales flush to zero below an amax of about 6 * 2^-10; an exact amax scale never does. + small = {"key": _fp8_pages((-8.5, -8.0), seed=5), "value": _fp8_pages((-8.5, -8.0), seed=6)} + identity, _ = _round_trip("identity", small) + tuned, _ = _round_trip("fp8_block_amax", small) + for role in small: + assert torch.count_nonzero(identity[role].float()) == 0 + assert torch.count_nonzero(tuned[role].float().reshape(-1, 16).abs().amax(-1)) == 2 * ( + _HEADS * _TOKENS * _HEAD_DIM // 16 + ) + + +def test_fp8_block_amax_second_host_trip_reproduces_the_first() -> None: + """A Page onboarded from host and evicted again is re-encoded from its decoded FP8 values.""" + + source = {"key": _fp8_pages((-2.0, 8.8), seed=7), "value": _fp8_pages((-2.0, 8.8), seed=8)} + first, codec = _round_trip("fp8_block_amax", source) + first_cold = [codec.cold_page(0), codec.cold_page(1)] + codec.encode([(2, 2), (3, 3)]) + assert torch.equal(codec.cold_page(2), first_cold[0]) + assert torch.equal(codec.cold_page(3), first_cold[1]) + codec.decode([(2, 0), (3, 1)]) + for role, buffer in codec.buffers.items(): + assert torch.equal(buffer[:2].view(torch.uint8), first[role].view(torch.uint8)), role diff --git a/tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py b/tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py index acc11208d3a8..97eb8b03e01c 100644 --- a/tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py +++ b/tests/unittest/_torch/kv_cache_compression/test_quantization_for_cold_page.py @@ -22,7 +22,11 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import DataType from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode from tensorrt_llm._torch.speculative.utils import update_spec_config_from_model_config -from tensorrt_llm.llmapi.llm_args import ColdPageQuantizationCompressionConfig, MTPDecodingConfig +from tensorrt_llm.llmapi.llm_args import ( + ColdPageQuantizationCompressionConfig, + DFlashDecodingConfig, + MTPDecodingConfig, +) from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization import QuantAlgo from tensorrt_llm.runtime.kv_cache_manager_v2 import ( @@ -41,6 +45,7 @@ def _manager( pretrained_config=None, skip_rope_quantization=False, nvfp4_residual_dim=64, + global_scale="identity", ): config = ColdPageQuantizationCompressionConfig( scale_checkpoint_path=( @@ -48,6 +53,7 @@ def _manager( ), skip_rope_quantization=skip_rope_quantization, nvfp4_residual_dim=nvfp4_residual_dim, + global_scale=global_scale, ) return Nvfp4ColdPageQuantizationCompression( config, @@ -201,8 +207,10 @@ def _deepseek_v4_csa_cache_config(*, num_model_layers=1, first_layer_id=0, eleme return SimpleNamespace(tokens_per_block=128, layers=tuple(layers)) -def _validate_compression(mode: object | None = None) -> None: - spec_config = None if mode is None else SimpleNamespace(spec_dec_mode=mode) +def _validate_compression(mode: object | None = None, *, spec_config: object | None = None) -> None: + assert mode is None or spec_config is None, "pass either mode or spec_config" + if spec_config is None and mode is not None: + spec_config = SimpleNamespace(spec_dec_mode=mode) with patch.object(util_mod, "is_sm_100f", return_value=True): util_mod.validate_kv_cache_compression_compatibility( ColdPageQuantizationCompressionConfig(), @@ -1285,12 +1293,41 @@ def test_runtime_admission_is_checked_before_manager_creation(monkeypatch) -> No def test_speculative_admission_accepts_verified_one_model_modes() -> None: _validate_compression(SpeculativeDecodingMode.EAGLE3_ONE_MODEL) _validate_compression(SpeculativeDecodingMode.MTP_EAGLE_ONE_MODEL) - for mode in ( - SpeculativeDecodingMode.MTP, - SpeculativeDecodingMode.DFLASH, - ): - with pytest.raises(ValueError, match="one-model MTP-EAGLE or EAGLE3"): - _validate_compression(mode) + with pytest.raises(ValueError, match="one-model MTP-EAGLE or EAGLE3"): + _validate_compression(SpeculativeDecodingMode.MTP) + + +@pytest.mark.parametrize("attention_backend", ["VANILLA", "FA4"]) +def test_speculative_admission_accepts_dflash_without_draft_kv_cache(attention_backend) -> None: + spec_config = DFlashDecodingConfig(max_draft_len=7, attention_backend=attention_backend) + assert spec_config.spec_dec_mode is SpeculativeDecodingMode.DFLASH + assert not util_mod.should_use_separate_draft_kv_cache(spec_config) + _validate_compression(spec_config=spec_config) + + +@pytest.mark.parametrize("enable_attention_dp", (False, True)) +@pytest.mark.parametrize("attention_backend", ["VANILLA", "FA4"]) +def test_admitted_dflash_builds_no_draft_kv_cache_and_adds_no_target_layers( + attention_backend, enable_attention_dp +) -> None: + """Admission implies the codec can only ever be handed the target KVCM.""" + + spec_config = DFlashDecodingConfig(max_draft_len=7, attention_backend=attention_backend) + _validate_compression(spec_config=spec_config) + creator = SimpleNamespace( + _speculative_config=spec_config, + _mapping=SimpleNamespace(enable_attention_dp=enable_attention_dp, pp_size=1), + _sparse_attention_config=None, + ) + assert not util_mod.KvCacheCreator._should_create_separate_draft_kv_cache(creator) + assert util_mod.get_num_spec_layers(spec_config) == 0 + + +def test_speculative_admission_rejects_dflash_with_draft_kv_cache() -> None: + spec_config = DFlashDecodingConfig(max_draft_len=7, attention_backend="TRTLLM") + assert util_mod.should_use_separate_draft_kv_cache(spec_config) + with pytest.raises(ValueError, match="DFlash without a separate draft KV cache"): + _validate_compression(spec_config=spec_config) def test_qwen35_mtp3_resolves_to_supported_one_model_mode() -> None: @@ -1774,3 +1811,185 @@ def test_measured_glm52_and_deepseek_v4_pages_are_reproduced() -> None: head_dim_per_layer=(), ) assert _configure_lifecycle(native, hot).cold_page_bytes == 384_512 + + +# --- global_scale ----------------------------------------------------------------------------- + + +def test_global_scale_defaults_to_identity_and_serializes() -> None: + config = ColdPageQuantizationCompressionConfig() + assert config.global_scale == "identity" + tuned = ColdPageQuantizationCompressionConfig(global_scale="fp8_block_amax") + assert ColdPageQuantizationCompressionConfig.model_validate(tuned.model_dump()) == tuned + with pytest.raises(ValueError, match="global_scale"): + ColdPageQuantizationCompressionConfig(global_scale="per_block") + + +def test_global_scale_and_scale_checkpoint_are_mutually_exclusive(tmp_path) -> None: + with pytest.raises(ValueError, match="mutually exclusive"): + ColdPageQuantizationCompressionConfig( + global_scale="fp8_block_amax", scale_checkpoint_path=str(tmp_path) + ) + + +def test_fp8_block_amax_global_scale_reaches_fp8_launch_metadata() -> None: + native, _ = _native() + (layout,) = _create( + _manager(global_scale="fp8_block_amax"), + _cache_config((0, "attention")), + native, + pp_layers=(0,), + num_kv_heads_per_layer=(4,), + head_dim_per_layer=(64,), + ) + assert _codec_state(native).runtime_type == 2 + assert [buffer.scales.nvfp4_orig_quant for buffer in layout.buffers] == [6.0, 6.0] + assert [buffer.scales.nvfp4_quant_orig for buffer in layout.buffers] == [1.0 / 6.0] * 2 + metadata = _configure_lifecycle(native, {0: {"key": 64 * 4 * 64, "value": 64 * 4 * 64}}) + expected = torch.tensor([6.0, 1.0 / 6.0, 1.0, 1.0], dtype=torch.float32) + assert torch.equal(metadata.scales[0], expected) + assert torch.equal(metadata.scales[1], expected) + # Geometry is unchanged: only the scale rows differ from the identity default. + identity_native, _ = _native() + _create( + _manager(), + _cache_config((0, "attention")), + identity_native, + pp_layers=(0,), + num_kv_heads_per_layer=(4,), + head_dim_per_layer=(64,), + ) + identity = _configure_lifecycle( + identity_native, {0: {"key": 64 * 4 * 64, "value": 64 * 4 * 64}} + ) + assert identity.cold_page_bytes == metadata.cold_page_bytes + assert torch.equal(identity.integers, metadata.integers) + assert identity.scales[0].tolist() == [1.0, 1.0, 1.0, 1.0] + + +def test_fp8_block_amax_global_scale_rejects_a_16_bit_target_cache() -> None: + native, _ = _native() + with pytest.raises(ValueError, match="requires an FP8 KV cache"): + _create( + _manager(global_scale="fp8_block_amax"), + _cache_config((0, "attention")), + native, + runtime_dtype=DataType.BF16, + pp_layers=(0,), + num_kv_heads_per_layer=(4,), + head_dim_per_layer=(64,), + ) + + +def test_fp8_block_amax_global_scale_keeps_identity_for_a_16_bit_draft_cache() -> None: + native, _ = _native() + (layout,) = _create( + _manager(global_scale="fp8_block_amax"), + _cache_config((0, "attention")), + native, + runtime_dtype=DataType.BF16, + pp_layers=(0,), + num_kv_heads_per_layer=(4,), + head_dim_per_layer=(64,), + is_draft=True, + ) + assert [buffer.scales.nvfp4_orig_quant for buffer in layout.buffers] == [1.0, 1.0] + + +def test_fp8_block_amax_global_scale_applies_to_deepseek_v4_compress_rows() -> None: + native, _ = _native() + manager = Nvfp4ColdPageQuantizationCompression( + ColdPageQuantizationCompressionConfig(global_scale="fp8_block_amax"), + pretrained_config=SimpleNamespace(model_type="deepseek_v4"), + ) + (layout,) = _create( + manager, + _deepseek_v4_csa_cache_config(element_bytes=1), + native, + pp_layers=(0,), + num_kv_heads_per_layer=(), + head_dim_per_layer=(), + ) + compressed = [buffer for buffer in layout.buffers if buffer.scales is not None] + assert compressed and all(buffer.scales.nvfp4_orig_quant == 6.0 for buffer in compressed) + + +# --- uncompressed_layer_types ------------------------------------------------------------------- + + +def _hybrid_window_cache_config(windows): + return SimpleNamespace( + tokens_per_block=32, + layers=tuple( + AttentionLayerConfig( + layer_id=layer_id, + buffers=[BufferConfig(role=role, size=4 * 32 * 64) for role in ("key", "value")], + sliding_window_size=window, + ) + for layer_id, window in enumerate(windows) + ), + ) + + +def _create_with_types(uncompressed_layer_types, windows): + native, _ = _native() + manager = Nvfp4ColdPageQuantizationCompression( + ColdPageQuantizationCompressionConfig(uncompressed_layer_types=uncompressed_layer_types), + pretrained_config=SimpleNamespace(model_type="muse_spark"), + ) + _create( + manager, + _hybrid_window_cache_config(windows), + native, + pp_layers=tuple(range(len(windows))), + num_kv_heads_per_layer=(4,) * len(windows), + head_dim_per_layer=(64,) * len(windows), + ) + return native, _codec_state(native) + + +@pytest.mark.parametrize( + ("uncompressed_layer_types", "expected_layer_ids"), + ( + ([], (0, 1, 2, 3)), + (["sliding_window"], (1, 3)), + (["full_attention"], (0, 2)), + (["sliding_window", "full_attention"], ()), + ), +) +def test_uncompressed_layer_types_leave_layers_to_the_lossless_codec( + uncompressed_layer_types, expected_layer_ids +) -> None: + _, state = _create_with_types(uncompressed_layer_types, (2048, None, 2048, None)) + assert state.layer_ids == expected_layer_ids + assert tuple(state.layer_layouts) == expected_layer_ids + + +def test_uncompressed_layer_types_keep_full_attention_lifecycle_quantized() -> None: + native, _ = _create_with_types(["sliding_window"], (2048, None, 2048, None)) + metadata = _configure_lifecycle( + native, {layer_id: {"key": 4 * 32 * 64, "value": 4 * 32 * 64} for layer_id in (1, 3)} + ) + assert metadata.num_buffers == 4 + assert metadata.integers[:4, 1].tolist() == [0, 0, 0, 0] + assert metadata.cold_page_bytes == 2 * (2 * 4 * 32 * 64 // 2 + 2 * 4 * 32 * 64 // 16) + + +def test_uniform_windows_count_as_full_attention() -> None: + _, state = _create_with_types(["sliding_window"], (2048, 2048)) + assert state.layer_ids == (0, 1) + _, state = _create_with_types(["full_attention"], (None, None)) + assert state.layer_ids == () + + +def test_uncompressed_layer_types_default_and_validation() -> None: + assert ColdPageQuantizationCompressionConfig().uncompressed_layer_types == [] + config = ColdPageQuantizationCompressionConfig(uncompressed_layer_types=["sliding_window"]) + assert ColdPageQuantizationCompressionConfig.model_validate(config.model_dump()) == config + with pytest.raises(ValueError, match="uncompressed_layer_types"): + ColdPageQuantizationCompressionConfig(uncompressed_layer_types=["mamba"]) + with pytest.raises(NotImplementedError, match="DeepSeek-V4"): + Nvfp4ColdPageQuantizationCompression( + ColdPageQuantizationCompressionConfig(uncompressed_layer_types=["sliding_window"]), + pretrained_config=SimpleNamespace(model_type="deepseek_v4"), + ) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 6758fdfea98f..625467a98543 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -4147,6 +4147,8 @@ def test_kv_cache_compression_config_dispatches_by_algorithm(): "skip_rope_quantization": False, "nvfp4_residual_dim": 64, "scale_checkpoint_path": None, + "global_scale": "identity", + "uncompressed_layer_types": [], } assert not cold_config.changes_physical_kv_length assert cold_config.supports_block_reuse() From 6ebd22e3e95bc8cd2d50156d3297be37b5de8aac Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:10:55 +0000 Subject: [PATCH 03/15] [None][feat] Disaggregated serving: opt-in generation KV prefetch admission and router held placement Generation workers: - TRTLLM_DISAGG_GEN_PREFETCH_REQUESTS=N (0 by default) admits up to N generation-only requests beyond max_num_active_requests, so their KV transfer overlaps the running decode. The forward batch stays capped at max_batch_size. The engine widens its sequence-slot pool and the KV cache manager v2 its lease budget by the same N. It is off under TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN, attention DP and pipeline parallelism. Requests that do not fit go back to the front of the waiting queue, which keeps FCFS order. Router: - LoadBalancingRouter gains max_inflight_per_server and hold_timeout_s. With a cap set, a request is placed only on a server below the cap, and otherwise held until one frees up; a request held longer than the timeout falls back to the least-loaded server. Without a cap, placement is unchanged. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../_torch/pyexecutor/disagg_gen_prefetch.py | 120 +++++ .../kv_cache/kv_cache_manager_v2.py | 9 + .../_torch/pyexecutor/model_engine.py | 28 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 50 ++- tensorrt_llm/serve/disagg_coordinator.py | 54 ++- tensorrt_llm/serve/router.py | 297 +++++++++++- .../kv_cache/test_kvcm2_integration.py | 65 ++- .../executor/test_disagg_gen_prefetch.py | 420 +++++++++++++++++ .../disaggregated/test_router_hold.py | 425 ++++++++++++++++++ 9 files changed, 1442 insertions(+), 26 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/disagg_gen_prefetch.py create mode 100644 tests/unittest/_torch/executor/test_disagg_gen_prefetch.py create mode 100644 tests/unittest/disaggregated/test_router_hold.py diff --git a/tensorrt_llm/_torch/pyexecutor/disagg_gen_prefetch.py b/tensorrt_llm/_torch/pyexecutor/disagg_gen_prefetch.py new file mode 100644 index 000000000000..35dae095ed50 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/disagg_gen_prefetch.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Generation-worker KV prefetch for disaggregated serving. + +A disaggregated generation worker normally admits at most ``max_batch_size`` +requests, so the next request's KV transfer only starts after a running +request finishes. Setting ``TRTLLM_DISAGG_GEN_PREFETCH_REQUESTS=N`` lets the +worker admit up to ``N`` more generation-only requests. They receive their KV +while the running batch decodes; the forward batch is still capped at +``max_batch_size`` by the scheduler budget. + +The extra requests need room in three pools, all widened by ``N``: + +* the request admission cap (``PyExecutor._get_request_admission_capacity``); +* the sequence-slot pool (``model_engine.max_num_seq_slots``). A prefetched + request is scheduled in the iteration where the request it replaces is + ``GENERATION_TO_COMPLETE`` and still owns its slot; +* the KV-cache index pool (``KVCacheManagerV2.max_admissible_sequences``). + +``N = 0`` (the default) leaves every pool and the admission cap unchanged. + +Prefetch is turned off when ``TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN=1``. +That scheduler option orders generation-only requests that have not emitted +their first token ahead of running decodes, so a prefetched request whose KV +has landed would take a running request's place in the forward batch, and with +it that request's speculative-decoding state, for an iteration. +""" + +import os + +from tensorrt_llm.logger import logger + +DISAGG_GEN_PREFETCH_REQUESTS_ENV = "TRTLLM_DISAGG_GEN_PREFETCH_REQUESTS" +# Read with the same rule as ``scheduler_v2`` so both agree on whether it is on. +PRIORITIZE_FIRST_TOKEN_ENV = "TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN" + + +def _prioritize_first_token_enabled() -> bool: + return os.environ.get(PRIORITIZE_FIRST_TOKEN_ENV, "0") == "1" + + +def get_disagg_gen_prefetch_requests() -> int: + """Return the requested number of prefetch requests (0 when unset).""" + raw = os.environ.get(DISAGG_GEN_PREFETCH_REQUESTS_ENV, "").strip() + if not raw: + return 0 + try: + value = int(raw) + except ValueError as e: + raise ValueError( + f"{DISAGG_GEN_PREFETCH_REQUESTS_ENV} must be a non-negative integer, got {raw!r}" + ) from e + if value < 0: + raise ValueError( + f"{DISAGG_GEN_PREFETCH_REQUESTS_ENV} must be a non-negative integer, got {value}" + ) + return value + + +def resolve_disagg_gen_prefetch_requests( + *, + is_disagg: bool, + enable_attention_dp: bool, + has_pp: bool, + is_hybrid: bool = False, + has_mrope_delta_cache: bool = False, + kv_cache_manager_is_v2: bool = True, + warn: bool = True, +) -> int: + """Return the effective number of prefetch requests for this executor. + + Prefetch is applied only where widening the sequence-slot pool is known + to be safe: a disaggregated executor with KV cache manager V2, no + attention DP, no pipeline parallelism, no pool that is sized from + ``max_batch_size`` instead of the sequence-slot pool (hybrid linear + attention, MRoPE delta cache), and a scheduler that never lets a request + that has not started decoding go ahead of a running one + (``TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN`` unset). Elsewhere the setting + is ignored. Every pool that prefetch widens is sized from this function, + so they agree on the value. + """ + requested = get_disagg_gen_prefetch_requests() + if requested == 0: + return 0 + reasons = [] + if not is_disagg: + reasons.append("no KV cache transceiver") + if enable_attention_dp: + reasons.append("attention DP") + if has_pp: + reasons.append("pipeline parallelism") + if is_hybrid: + reasons.append("hybrid linear attention") + if has_mrope_delta_cache: + reasons.append("MRoPE delta cache") + if not kv_cache_manager_is_v2: + reasons.append("KV cache manager V1") + if _prioritize_first_token_enabled(): + reasons.append(f"{PRIORITIZE_FIRST_TOKEN_ENV}=1") + if reasons: + if warn: + logger.warning_once( + f"{DISAGG_GEN_PREFETCH_REQUESTS_ENV}={requested} is ignored " + f"({', '.join(reasons)}).", + key="disagg_gen_prefetch_ignored", + ) + return 0 + return requested diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 22fccbabf348..2eb972175fbc 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -109,6 +109,7 @@ from ..config_utils import get_layer_attention_window, uses_vswa_kv_cache_layout from ..connectors.kv_cache_connector import KvCacheConnectorManager +from ..disagg_gen_prefetch import resolve_disagg_gen_prefetch_requests from ..kv_cache_events import StreamingKVCacheEventManager, validate_streaming_support from ..kv_cache_stats import ( KVCacheV2IterationStatsReport, @@ -2228,6 +2229,14 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self._fresh_fill_announced = False self._fresh_fill_unavailable_announced = False self.max_admissible_sequences = max_num_sequences * (2 if needs_extra_leases else 1) + # A disaggregated generation worker that prefetches KV for requests + # beyond max_batch_size holds that many more leases at once. + self.max_admissible_sequences += resolve_disagg_gen_prefetch_requests( + is_disagg=is_disagg, + enable_attention_dp=mapping.enable_attention_dp, + has_pp=mapping.has_pp(), + warn=False, + ) # The guard page is held by a permanent sequence, so it needs an index # slot of its own. Taking one of the scheduler's would change which # requests get admitted, so a run with the diagnostic on would no diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5c74421ea255..8fded4eed7fd 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -46,6 +46,7 @@ from ..utils import get_per_request_prefill_cuda_graph_flag, set_torch_compiling from .config_utils import is_hybrid_linear from .cuda_graph_runner import CUDAGraphRunner +from .disagg_gen_prefetch import resolve_disagg_gen_prefetch_requests from .engine.cuda_graph import filter_cuda_graph_batch_sizes from .engine.lora import make_lora_model_config from .engine.model_call import ModelCaller @@ -262,7 +263,7 @@ def __init__( self.mapping = mapping if mapping.has_pp(): init_pp_comm(mapping) - from ._util import (compute_max_num_sequences, + from ._util import (compute_max_num_sequences, is_disagg_enabled, resolved_kv_cache_manager_is_v2, should_enable_adp_dummy_fixes, should_enable_non_overlap_adp_forward_intent, @@ -363,20 +364,33 @@ def __init__( self._enable_non_overlap_adp_forward_intent = ( should_enable_non_overlap_adp_forward_intent( mapping, llm_args.disable_overlap_scheduler)) + kv_cache_manager_is_v2 = resolved_kv_cache_manager_is_v2( + llm_args.kv_cache_config, self.max_beam_width) + is_hybrid = is_hybrid_linear(pretrained_config) + has_mrope_delta_cache = resolve_mrope_position_deltas_cache( + self.model) is not None self._enable_overlap_headroom = should_enable_overlap_headroom( mapping, llm_args.disable_overlap_scheduler, - kv_cache_manager_is_v2=resolved_kv_cache_manager_is_v2( - llm_args.kv_cache_config, self.max_beam_width), - is_hybrid=is_hybrid_linear(pretrained_config), - has_mrope_delta_cache=resolve_mrope_position_deltas_cache( - self.model) is not None) + kv_cache_manager_is_v2=kv_cache_manager_is_v2, + is_hybrid=is_hybrid, + has_mrope_delta_cache=has_mrope_delta_cache) + # Seats for the generation requests a disaggregated worker admits + # beyond max_batch_size to prefetch their KV (0 unless opted in). + self.disagg_gen_prefetch_requests = resolve_disagg_gen_prefetch_requests( + is_disagg=is_disagg_enabled( + getattr(llm_args, "cache_transceiver_config", None)), + enable_attention_dp=mapping.enable_attention_dp, + has_pp=mapping.has_pp(), + is_hybrid=is_hybrid, + has_mrope_delta_cache=has_mrope_delta_cache, + kv_cache_manager_is_v2=kv_cache_manager_is_v2) self.max_num_seq_slots = compute_max_num_sequences( mapping, self.batch_size, llm_args.disable_overlap_scheduler, enable_overlap_headroom=self._enable_overlap_headroom, - ) + ) + self.disagg_gen_prefetch_requests self.sparse_attention_config = self.model.model_config.sparse_attention_config # In case that some tests use stub models and override `_load_model`. if not hasattr(self.model, 'extra_attrs'): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 3d18f15bbaec..929b93ac677e 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -732,6 +732,19 @@ def __init__( f"{required_max_num_tokens} to run at the declared " f"max_batch_size.") self.max_num_active_requests = derived_cap + # Generation-only requests admitted beyond max_num_active_requests so + # their KV transfer overlaps the running decode. The engine widened + # its sequence-slot pool by the same amount; 0 unless opted in. + self.disagg_gen_prefetch_requests = (getattr( + model_engine, "disagg_gen_prefetch_requests", + 0) if kv_cache_transceiver is not None else 0) + if self.disagg_gen_prefetch_requests > 0: + logger.info( + f"[PyExecutor] disaggregated generation KV prefetch: admitting " + f"up to {self.disagg_gen_prefetch_requests} generation-only " + f"requests beyond max_num_active_requests=" + f"{self.max_num_active_requests}; the forward batch is still " + f"capped at max_batch_size") self.active_requests: List[LlmRequest] = [] self.expected_num_active_requests = 0 # Buffer for responses staged by rank-divergent paths (disagg transfer @@ -6176,7 +6189,38 @@ def _get_request_admission_capacity(self) -> int: """Return the maximum number of distinct requests admitted at once.""" if self.enable_attention_dp: return self.dist.tp_size * self.max_num_active_requests - return self.max_num_active_requests + return (self.max_num_active_requests + + getattr(self, "disagg_gen_prefetch_requests", 0)) + + def _limit_prefetch_admission( + self, waiting_queue: WaitingQueue, + new_requests: List[RequestQueueItem], + total_num_active_requests: int) -> List[RequestQueueItem]: + """Admit only generation-only requests into the KV prefetch slots. + + Popped items (with their child requests) that fit within + ``max_num_active_requests - total_num_active_requests`` use the + regular slots. Later items use the prefetch slots and must be + generation-only requests without child requests; the first item that + is not, and everything after it, go back to the front of the waiting + queue so FCFS order is kept. + """ + regular = self.max_num_active_requests - total_num_active_requests + cut = len(new_requests) + used = 0 + for idx, item in enumerate(new_requests): + used += 1 + len(item.child_req_ids or ()) + if used <= regular: + continue + request = item.request + if (item.child_req_ids or request is None + or getattr(request, "request_type", None) + != RequestType.REQUEST_TYPE_GENERATION_ONLY): + cut = idx + break + if cut < len(new_requests): + waiting_queue.prepend_requests(new_requests[cut:]) + return new_requests[:cut] def _configure_benchmark_req_queues_size(self) -> None: """Clamp the benchmark fill target to the request admission ceiling.""" @@ -6296,6 +6340,10 @@ def _pop_from_waiting_queue( enable_attention_dp=self.enable_attention_dp, max_num_active_requests=self.max_num_active_requests, all_ranks_num_active_requests=all_ranks_num_active_requests) + if (new_requests and not self.enable_attention_dp + and getattr(self, "disagg_gen_prefetch_requests", 0) > 0): + new_requests = self._limit_prefetch_admission( + waiting_queue, new_requests, total_num_active_requests) if (not new_requests or not getattr(self, "_mm_encoder_item_scheduling_enabled", False) or self.enable_attention_dp): diff --git a/tensorrt_llm/serve/disagg_coordinator.py b/tensorrt_llm/serve/disagg_coordinator.py index a396d4f90213..4dda8c510827 100644 --- a/tensorrt_llm/serve/disagg_coordinator.py +++ b/tensorrt_llm/serve/disagg_coordinator.py @@ -227,19 +227,57 @@ async def select( normalized_role = self._normalize_role(role) router = self._router_for_role(normalized_role) reservation_key = (normalized_role, req_id) + if getattr(router, "placement_hold_timeout_s", None) is not None: + server, info, request_id = await self._select_with_hold( + router, reservation_key, routing_key, req_id, exclude_server + ) + else: + async with self._reservation_lock(reservation_key): + previous = self._reservation_tasks.pop(reservation_key, None) + if previous is not None: + previous.cancel() + await router.finish_request_by_id(req_id, False) + server, info, request_id = await router.get_next_server_by_key( + routing_key, req_id=req_id, exclude_server=exclude_server + ) + self._reservation_tasks[reservation_key] = asyncio.create_task( + self._expire_reservation(reservation_key, router) + ) + self._api_lat(f"select[{role}]").record(time.monotonic() - _t0) + return server, self._compact_route_info(info), request_id + + async def _select_with_hold( + self, + router: Router, + reservation_key: tuple[str, int], + routing_key, + req_id, + exclude_server: Optional[str], + ) -> Tuple[str, dict, Optional[str]]: + """Place through a router that may hold the request until a server frees. + + The placement runs outside the reservation lock: the lock stripes are + shared with ``finish``, and a finish is what frees the server a held + placement waits for. + """ async with self._reservation_lock(reservation_key): previous = self._reservation_tasks.pop(reservation_key, None) if previous is not None: previous.cancel() await router.finish_request_by_id(req_id, False) - server, info, request_id = await router.get_next_server_by_key( - routing_key, req_id=req_id, exclude_server=exclude_server - ) - self._reservation_tasks[reservation_key] = asyncio.create_task( - self._expire_reservation(reservation_key, router) - ) - self._api_lat(f"select[{role}]").record(time.monotonic() - _t0) - return server, self._compact_route_info(info), request_id + server, info, request_id = await router.get_next_server_by_key( + routing_key, req_id=req_id, exclude_server=exclude_server + ) + async with self._reservation_lock(reservation_key): + stale = self._reservation_tasks.pop(reservation_key, None) + if stale is not None: + stale.cancel() + # A finish that landed between placement and here already released it. + if router.has_placement(req_id): + self._reservation_tasks[reservation_key] = asyncio.create_task( + self._expire_reservation(reservation_key, router) + ) + return server, info, request_id @staticmethod def _compact_route_info(info: dict) -> dict: diff --git a/tensorrt_llm/serve/router.py b/tensorrt_llm/serve/router.py index 03524723b5c4..c836b0089856 100644 --- a/tensorrt_llm/serve/router.py +++ b/tensorrt_llm/serve/router.py @@ -13,9 +13,11 @@ # limitations under the License. import asyncio +import math import time from abc import ABC, abstractmethod from collections import OrderedDict +from dataclasses import dataclass from typing import Awaitable, Callable, Dict, Iterable, List, Optional import aiohttp @@ -53,6 +55,27 @@ # requests (placeholder message content) still key per conversation. ROUTE_AFFINITY_TOKEN_PREFIX = 256 +# Held placement (LoadBalancingRouter ``max_inflight_per_server``): default +# longest wait before a held request falls back to the least-loaded server. +PLACEMENT_HOLD_TIMEOUT_DEFAULT_S = 300.0 +# Request IDs released before their placement arrived are remembered this long +# (and at most this many) so the late placement is refused instead of leaking. +PLACEMENT_EARLY_RELEASE_TTL_S = 600.0 +PLACEMENT_EARLY_RELEASE_MAX = 16384 + + +class HeldPlacementCancelled(ValueError): + """A held placement was released before a server freed up.""" + + +@dataclass +class _HeldPlacement: + """One request waiting in the router for a server below the in-flight cap.""" + exclude_server: Optional[str] + num_tokens: int + enqueued_at: float + future: asyncio.Future + class ServerState: @@ -93,6 +116,11 @@ async def decrement_load_by(self, num_tokens: int = 0): self._num_active_requests -= 1 self._num_active_tokens -= num_tokens + def adjust_load_nowait(self, num_requests: int, num_tokens: int) -> None: + """Adjust the load without awaiting, for callers that must not yield.""" + self._num_active_requests += num_requests + self._num_active_tokens += num_tokens + async def is_healthy(self) -> bool: try: async with self._session.get( @@ -752,6 +780,21 @@ async def finish_request(self, class LoadBalancingRouter(LoadBalancingMixin, Router): + """Join-shortest-queue router over requests (or prompt tokens). + + With ``max_inflight_per_server`` set, a request is placed only on a server + with fewer than that many requests in flight. When every server is at the + cap the request is held in the router, first come first served, until a + finish frees a slot. It then goes to the server with the fewest requests + in flight; among equally busy servers, to the one whose oldest in-flight + request started last. A request held longer than ``hold_timeout_s`` + falls back to the least-loaded server. Without the cap, placement is + immediate. + + The held path never awaits while it changes routing state, so it does not + take ``self._lock``: each of its state changes is atomic with respect to + the event loop. + """ def __init__(self, server_role: ServerRole, @@ -759,6 +802,8 @@ def __init__(self, metadata_server_cfg: MetadataServerConfig = None, metadata_server: JsonDictionary = None, use_tokens: bool = False, + max_inflight_per_server: Optional[int] = None, + hold_timeout_s: float = PLACEMENT_HOLD_TIMEOUT_DEFAULT_S, **kwargs): super().__init__(server_role, servers, metadata_server_cfg, metadata_server, **kwargs) @@ -766,6 +811,48 @@ def __init__(self, # The coordinator path cannot retain the request object, so preserve its # token weight by disaggregated request ID until /finish arrives. self._coord_request_tokens: dict[int, int] = {} + if max_inflight_per_server is not None and ( + isinstance(max_inflight_per_server, bool) + or not isinstance(max_inflight_per_server, int) + or max_inflight_per_server < 1): + raise ValueError("max_inflight_per_server must be a positive " + f"integer, got {max_inflight_per_server!r}") + hold_timeout_s = float(hold_timeout_s) + if not hold_timeout_s > 0: + raise ValueError( + f"hold_timeout_s must be positive, got {hold_timeout_s}") + self._max_inflight_per_server = max_inflight_per_server + self._hold_timeout_s = hold_timeout_s + # Held requests in arrival order, keyed like _req_routing_table. + self._held: OrderedDict[int, _HeldPlacement] = OrderedDict() + # Token weight and start time of each placement made under the cap. + self._hold_tokens: dict[int, int] = {} + self._hold_started: dict[int, float] = {} + self._early_releases: OrderedDict[int, float] = OrderedDict() + self._hold_fallbacks = 0 + self._hold_wait_lat = None + if max_inflight_per_server is not None: + from tensorrt_llm.serve.responses_utils import PeriodicLatencyLogger + self._hold_wait_lat = PeriodicLatencyLogger( + f"router.hold_wait[{server_role}]") + logger.info( + f"LoadBalancingRouter[{server_role}]: holding requests while " + f"every server has >= {max_inflight_per_server} in flight " + f"(hold_timeout_s={hold_timeout_s})") + + @property + def placement_hold_timeout_s(self) -> Optional[float]: + """Longest time a placement may wait, or None if it never waits.""" + if self._max_inflight_per_server is None: + return None + return self._hold_timeout_s + + @property + def num_held_requests(self) -> int: + return len(self._held) + + def has_placement(self, key: int) -> bool: + return key in self._req_routing_table def _on_servers_updated(self, old_servers, new_servers): new_state = {} @@ -773,12 +860,160 @@ def _on_servers_updated(self, old_servers, new_servers): new_state[server] = (self._server_state.get(server) or self._create_server_state(server)) self._server_state = new_state + if self._max_inflight_per_server is not None: + self._dispatch_held() + + # ---- held placement (max_inflight_per_server) ------------------------------- + + def _hold_pick_server(self, exclude_server: Optional[str]) -> Optional[str]: + """Return a server below the in-flight cap, or None if all are full.""" + cap = self._max_inflight_per_server + eligible = [ + server for server in self._servers + if server != exclude_server and server in self._server_state + and self._server_state[server]._num_active_requests < cap + ] + if not eligible: + return None + oldest_start: dict[str, float] = {} + for key, started in self._hold_started.items(): + server = self._req_routing_table.get(key) + if server is not None and started < oldest_start.get( + server, math.inf): + oldest_start[server] = started + ranks = { + server: + (self._server_state[server]._num_active_requests, + -oldest_start.get(server, math.inf), self._get_server_load(server)) + for server in eligible + } + best = min(ranks.values()) + tied = [server for server in eligible if ranks[server] == best] + server = tied[self._rr_counter % len(tied)] + self._rr_counter += 1 + return server + + def _hold_commit(self, server: str, key: int, num_tokens: int) -> None: + if key in self._req_routing_table: + # A repeated placement for the same key replaces the old one. + self._hold_release(key, dispatch=False) + self._server_state[server].adjust_load_nowait(1, num_tokens) + self._req_routing_table[key] = server + self._hold_tokens[key] = num_tokens + self._hold_started[key] = time.monotonic() + + def _hold_release(self, key: int, dispatch: bool = True) -> bool: + server = self._req_routing_table.pop(key, None) + num_tokens = self._hold_tokens.pop(key, 0) + self._hold_started.pop(key, None) + if server is None: + return False + if server in self._server_state: + self._server_state[server].adjust_load_nowait(-1, -num_tokens) + if dispatch: + self._dispatch_held() + return True + + def _dispatch_held(self) -> None: + """Place held requests, oldest first, while servers have room.""" + for key, waiter in list(self._held.items()): + if waiter.future.done(): + del self._held[key] + continue + server = self._hold_pick_server(waiter.exclude_server) + if server is None: + continue + del self._held[key] + self._hold_commit(server, key, waiter.num_tokens) + waiter.future.set_result(server) + + def _remember_early_release(self, key: int) -> None: + now = time.monotonic() + self._early_releases[key] = now + self._early_releases.move_to_end(key) + while self._early_releases and ( + len(self._early_releases) > PLACEMENT_EARLY_RELEASE_MAX + or now - next(iter(self._early_releases.values())) + > PLACEMENT_EARLY_RELEASE_TTL_S): + self._early_releases.popitem(last=False) + + async def _place_with_hold(self, key: int, exclude_server: Optional[str], + num_tokens: int) -> str: + self._validate_servers_available() + previous = self._held.pop(key, None) + if previous is not None and not previous.future.done(): + previous.future.set_exception( + HeldPlacementCancelled( + f"held placement {key} superseded by a new placement")) + waiter = _HeldPlacement( + exclude_server=exclude_server, + num_tokens=num_tokens, + enqueued_at=time.monotonic(), + future=asyncio.get_running_loop().create_future()) + self._held[key] = waiter + # Older held requests are served first; this one is placed now only if + # a server still has room after them. + self._dispatch_held() + if waiter.future.done(): + return waiter.future.result() + try: + server = await asyncio.wait_for(asyncio.shield(waiter.future), + timeout=self._hold_timeout_s) + except asyncio.TimeoutError: + server = self._place_after_hold_timeout(key, waiter) + except asyncio.CancelledError: + self._abandon_held(key, waiter) + raise + self._hold_wait_lat.record(time.monotonic() - waiter.enqueued_at) + return server + + def _place_after_hold_timeout(self, key: int, + waiter: _HeldPlacement) -> str: + if waiter.future.done(): + # Placed or released in the same loop turn as the timeout. + return waiter.future.result() + if self._held.get(key) is waiter: + del self._held[key] + waiter.future.cancel() + server = self._select_least_loaded(waiter.exclude_server) + if server is None: + raise ValueError(f"No available servers after excluding " + f"{waiter.exclude_server}") + self._hold_commit(server, key, waiter.num_tokens) + self._hold_fallbacks += 1 + if self._hold_fallbacks == 1 or self._hold_fallbacks % 100 == 0: + logger.warning( + f"LoadBalancingRouter[{self._server_role}]: request held " + f"{self._hold_timeout_s}s without a server below " + f"max_inflight_per_server={self._max_inflight_per_server}; " + f"placed on least-loaded {server} " + f"(fallbacks={self._hold_fallbacks})") + return server + + def _abandon_held(self, key: int, waiter: _HeldPlacement) -> None: + """Drop a held placement whose caller was cancelled.""" + if self._held.get(key) is waiter: + del self._held[key] + future = waiter.future + if not future.done(): + future.cancel() + elif (not future.cancelled() and future.exception() is None + and self._req_routing_table.get(key) == future.result()): + # A server was assigned just before the cancellation landed. + self._hold_release(key) async def get_next_server(self, request: OpenAIRequest, exclude_server: Optional[str] = None, req_id: Optional[int] = None) -> tuple[str, dict]: del req_id + if self._max_inflight_per_server is not None: + num_tokens = (get_request_num_tokens(request) + if self._use_tokens else 0) + server = await self._place_with_hold(id(request), exclude_server, + num_tokens) + return server, {"server_info": self._server_info.get(server, {})} + self._validate_servers_available() async with self._lock: @@ -796,6 +1031,9 @@ async def finish_request(self, success: bool = True, req_id: Optional[int] = None): del session, success, req_id + if self._max_inflight_per_server is not None: + self._hold_release(id(request)) + return async with self._lock: await self._unregister_request(request) @@ -816,8 +1054,17 @@ async def get_next_server_by_key(self, if req_id is None: raise ValueError( "Load-balancing coordinator routing requires a request ID") - self._validate_servers_available() num_tokens = int((routing_key or {}).get("num_tokens", 0)) + if self._max_inflight_per_server is not None: + if self._early_releases.pop(req_id, None) is not None: + raise HeldPlacementCancelled( + f"request {req_id} was released before its placement") + server = await self._place_with_hold(req_id, exclude_server, + num_tokens) + return server, { + "server_info": self._server_info.get(server, {}) + }, req_id + self._validate_servers_available() async with self._lock: server = self._select_least_loaded(exclude_server) if server is None: @@ -834,6 +1081,17 @@ async def finish_request_by_id(self, req_id, success=True): del success if req_id is None: return + if self._max_inflight_per_server is not None: + waiter = self._held.pop(req_id, None) + if waiter is not None: + if not waiter.future.done(): + waiter.future.set_exception( + HeldPlacementCancelled( + f"request {req_id} was released while held")) + return + if not self._hold_release(req_id): + self._remember_early_release(req_id) + return async with self._lock: server = self._req_routing_table.pop(req_id, None) num_tokens = self._coord_request_tokens.pop(req_id, 0) @@ -1823,12 +2081,30 @@ async def get_next_server(self, "exclude_server": exclude_server } _t0 = time.monotonic() - body = await self._post_select(payload) + hold_timeout_s = getattr(self._local, "placement_hold_timeout_s", None) + if hold_timeout_s is None: + body = await self._post_select(payload) + else: + # The coordinator may hold this placement for up to hold_timeout_s + # and then places it anyway, so wait that long plus the usual + # request timeout. A /select given up here is still held there + # (the coordinator handler is not cancelled by a disconnect), so + # release it explicitly. + try: + body = await self._post_select(payload, + timeout=hold_timeout_s + + self._request_timeout_s) + except (asyncio.CancelledError, asyncio.TimeoutError, + aiohttp.ClientError): + self._enqueue_finish(payload["req_id"], False) + raise self._select_lat.record(time.monotonic() - _t0) info = body.get("info") or {} return body["server"], info - async def _post_select(self, payload: dict) -> dict: + async def _post_select(self, + payload: dict, + timeout: Optional[float] = None) -> dict: """POST /select, retrying connection-level failures. A pooled connection the coordinator closed first fails instantly with @@ -1841,13 +2117,14 @@ async def _post_select(self, payload: dict) -> dict: is already overloaded, so retrying would only add load. """ data = msgpack.packb(payload, use_bin_type=True) + if timeout is None: + timeout = self._request_timeout_s for attempt in range(1, COORDINATOR_SELECT_MAX_ATTEMPTS + 1): try: - async with self.session.post( - f"{self._coordinator_url}/select", - data=data, - headers=_MSGPACK_HEADERS, - timeout=self._request_timeout_s) as resp: + async with self.session.post(f"{self._coordinator_url}/select", + data=data, + headers=_MSGPACK_HEADERS, + timeout=timeout) as resp: body = msgpack.unpackb(await resp.read(), raw=False) if resp.status != 200: raise ValueError( @@ -1872,7 +2149,9 @@ async def finish_request(self, # request path. A fixed worker pool bounds tasks and connections during a # coordinator outage; overflow relies on coordinator-side expiration. del session - req_id = self._request_id(request, req_id) + self._enqueue_finish(self._request_id(request, req_id), success) + + def _enqueue_finish(self, req_id: int, success: bool) -> None: self._ensure_finish_workers() try: self._finish_queue.put_nowait((req_id, success)) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index dcc0110b49a3..a09ef52c8149 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -3069,6 +3069,8 @@ def _index_mapper_capacity_for( is_disagg: bool = False, num_reserved_index_slots: int = 1, disable_overlap_scheduler: bool = True, + prefetch_requests: str = "", + prioritize_first_token: str = "", ) -> tuple[int, int, int]: """Construct a manager and return the three sizes it derives. @@ -3079,6 +3081,11 @@ def _index_mapper_capacity_for( would be an out-of-bounds write. The third is the published admission bound that ``validate_seq_slot_pool_covers_admission`` checks the seat pool against at startup, and it is the capacity minus the reserved dummy slots. + + ``prefetch_requests`` and ``prioritize_first_token`` are the values of + ``TRTLLM_DISAGG_GEN_PREFETCH_REQUESTS`` and + ``TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN`` during construction; both are + pinned (unset by default) so an exported value cannot change the result. """ module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2" fake_impl = Mock() @@ -3102,7 +3109,12 @@ def build_base_config( # Pin both diagnostics off: the guard page would add an index slot. patch.dict( os.environ, - {"TRTLLM_KV_GUARD_PAGE": "", "TRTLLM_KV_FRESH_PAGE_FILL": ""}, + { + "TRTLLM_KV_GUARD_PAGE": "", + "TRTLLM_KV_FRESH_PAGE_FILL": "", + "TRTLLM_DISAGG_GEN_PREFETCH_REQUESTS": prefetch_requests, + "TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN": prioritize_first_token, + }, ), patch(f"{module}.IndexMapper") as index_mapper_cls, patch(f"{module}.KVCacheManagerPy", Mock(return_value=fake_impl)), @@ -3190,6 +3202,57 @@ def test_index_mapper_capacity_covers_the_overlapping_cohorts( assert max_admissible_sequences == expected - reserved +# (max_batch_size, pp_size, disable_overlap_scheduler, is_disagg, prefetch, +# prioritize_first_token, expected capacity) with one reserved slot. +# +# TRTLLM_DISAGG_GEN_PREFETCH_REQUESTS=N adds N leases on top of the formula +# above, for a non-PP disaggregated manager only. +_PREFETCH_INDEX_MAPPER_CAPACITY_CASES = [ + pytest.param(2, 1, False, True, "1", "", 6, id="disagg_overlap_on_n1"), + pytest.param(2, 1, True, True, "2", "", 7, id="disagg_overlap_off_n2"), + pytest.param(1, 1, False, True, "1", "", 4, id="disagg_b1_n1"), + pytest.param(2, 1, False, True, "0", "", 5, id="disagg_n0_unchanged"), + pytest.param(2, 1, False, False, "1", "", 5, id="aggregated_unchanged"), + pytest.param(2, 4, True, True, "1", "", 17, id="pp4_disagg_unchanged"), + # The scheduler option that lets a first-token request go ahead of a + # running decode turns prefetch off, here as in the seat pool. + pytest.param(2, 1, False, True, "1", "1", 5, id="prioritize_first_token_unchanged"), +] + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "max_batch_size,pp_size,disable_overlap_scheduler,is_disagg,prefetch,prioritize,expected", + _PREFETCH_INDEX_MAPPER_CAPACITY_CASES, +) +def test_index_mapper_capacity_adds_the_disagg_gen_prefetch_leases( + max_batch_size: int, + pp_size: int, + disable_overlap_scheduler: bool, + is_disagg: bool, + prefetch: str, + prioritize: str, + expected: int, +) -> None: + capacity, page_table_capacity, max_admissible_sequences = _index_mapper_capacity_for( + max_batch_size=max_batch_size, + pp_size=pp_size, + is_disagg=is_disagg, + disable_overlap_scheduler=disable_overlap_scheduler, + prefetch_requests=prefetch, + prioritize_first_token=prioritize, + ) + assert capacity == expected + assert page_table_capacity == expected + assert max_admissible_sequences == expected - 1 + + +@pytest.mark.cpu_only +def test_index_mapper_rejects_an_invalid_prefetch_setting() -> None: + with pytest.raises(ValueError, match="TRTLLM_DISAGG_GEN_PREFETCH_REQUESTS"): + _index_mapper_capacity_for(max_batch_size=2, is_disagg=True, prefetch_requests="-1") + + @pytest.mark.parametrize( "setting, expected", [ diff --git a/tests/unittest/_torch/executor/test_disagg_gen_prefetch.py b/tests/unittest/_torch/executor/test_disagg_gen_prefetch.py new file mode 100644 index 000000000000..769ff131eb42 --- /dev/null +++ b/tests/unittest/_torch/executor/test_disagg_gen_prefetch.py @@ -0,0 +1,420 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Disaggregated generation KV prefetch (TRTLLM_DISAGG_GEN_PREFETCH_REQUESTS). + +The KV index pool part (``KVCacheManagerV2.max_admissible_sequences``) is +covered next to the existing capacity formula in +``kv_cache/test_kvcm2_integration.py``. +""" + +from collections import deque +from types import MethodType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +import tensorrt_llm._torch.pyexecutor.model_engine as model_engine_module +import tensorrt_llm._torch.speculative.dflash as dflash_module +from tensorrt_llm._torch.pyexecutor import _util +from tensorrt_llm._torch.pyexecutor.disagg_gen_prefetch import ( + DISAGG_GEN_PREFETCH_REQUESTS_ENV, + PRIORITIZE_FIRST_TOKEN_ENV, + get_disagg_gen_prefetch_requests, + resolve_disagg_gen_prefetch_requests, +) +from tensorrt_llm._torch.pyexecutor.executor_request_queue import RequestQueueItem +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.resource_manager import NoFreeSlotsError +from tensorrt_llm._torch.pyexecutor.scheduler.waiting_queue import FCFSWaitingQueue +from tensorrt_llm._torch.pyexecutor.seq_slot_manager import SeqSlotManager +from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata, DFlashWorker +from tensorrt_llm.bindings.executor import RequestType +from tensorrt_llm.executor.request import DEFAULT_REQUEST_PRIORITY +from tensorrt_llm.llmapi import KvCacheConfig +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.cpu_only + +_ELIGIBLE = dict(is_disagg=True, enable_attention_dp=False, has_pp=False) + + +@pytest.fixture(autouse=True) +def _scheduler_option_unset(monkeypatch): + monkeypatch.delenv(PRIORITIZE_FIRST_TOKEN_ENV, raising=False) + + +@pytest.fixture +def prefetch_env(monkeypatch): + def _set(value): + if value is None: + monkeypatch.delenv(DISAGG_GEN_PREFETCH_REQUESTS_ENV, raising=False) + else: + monkeypatch.setenv(DISAGG_GEN_PREFETCH_REQUESTS_ENV, value) + + return _set + + +@pytest.mark.parametrize("raw, expected", [(None, 0), ("", 0), ("0", 0), ("1", 1), (" 3 ", 3)]) +def test_env_parsing(prefetch_env, raw, expected): + prefetch_env(raw) + assert get_disagg_gen_prefetch_requests() == expected + + +@pytest.mark.parametrize("raw", ["-1", "one", "1.5"]) +def test_env_rejects_invalid_values(prefetch_env, raw): + prefetch_env(raw) + with pytest.raises(ValueError): + get_disagg_gen_prefetch_requests() + + +@pytest.mark.parametrize( + "override", + [ + {"is_disagg": False}, + {"enable_attention_dp": True}, + {"has_pp": True}, + {"is_hybrid": True}, + {"has_mrope_delta_cache": True}, + {"kv_cache_manager_is_v2": False}, + ], +) +def test_resolve_ignores_unsupported_executors(prefetch_env, override): + prefetch_env("2") + kwargs = {**_ELIGIBLE, **override} + assert resolve_disagg_gen_prefetch_requests(**kwargs) == 0 + + +def test_resolve_gate(prefetch_env): + prefetch_env(None) + assert resolve_disagg_gen_prefetch_requests(**_ELIGIBLE) == 0 + prefetch_env("2") + assert resolve_disagg_gen_prefetch_requests(**_ELIGIBLE) == 2 + + +@pytest.mark.parametrize("value, prefetch", [("1", 0), ("0", 2), ("true", 2)]) +def test_resolve_follows_the_prioritize_first_token_scheduler_option( + prefetch_env, monkeypatch, value, prefetch +): + # The scheduler turns the option on only for exactly "1"; prefetch is + # turned off in exactly the same cases. + prefetch_env("2") + monkeypatch.setenv(PRIORITIZE_FIRST_TOKEN_ENV, value) + assert resolve_disagg_gen_prefetch_requests(**_ELIGIBLE) == prefetch + + +def _engine_seat_pool(monkeypatch, *, cache_transceiver_config, max_batch_size=2): + """Run the real ``PyTorchModelEngine.__init__`` up to the seat-pool sizing. + + Returns ``(max_num_seq_slots, disagg_gen_prefetch_requests)``. The gates + that need a real model are pinned to an aggregated-style non-ADP, non-PP + V2 engine; ``compute_max_num_sequences`` is the real one, so the pool is + ``max_batch_size`` plus whatever the prefetch term adds. + """ + for name in ( + "should_enable_adp_dummy_fixes", + "should_enable_scheduler_aware_adp_dummy", + "should_enable_non_overlap_adp_forward_intent", + "should_enable_overlap_headroom", + ): + monkeypatch.setattr(_util, name, Mock(return_value=False)) + monkeypatch.setattr(_util, "resolved_kv_cache_manager_is_v2", Mock(return_value=True)) + for name in ("_configure_deep_gemm_pdl", "create_input_processor"): + monkeypatch.setattr(model_engine_module, name, Mock()) + monkeypatch.setattr( + model_engine_module, "resolve_mrope_position_deltas_cache", lambda model: None + ) + monkeypatch.setattr(model_engine_module, "is_hybrid_linear", lambda config: False) + monkeypatch.setattr(PyTorchModelEngine, "_validate_breakable_cuda_graph_compatibility", Mock()) + monkeypatch.setattr(PyTorchModelEngine, "__del__", lambda self: None) + # The seat pool is sized before multimodal scheduling is set up; stop there. + monkeypatch.setattr( + model_engine_module.MultimodalItemScheduler, + "maybe_create", + Mock(side_effect=RuntimeError("seat pool sized")), + ) + model = torch.nn.Module() + model.model_config = SimpleNamespace( + pretrained_config=SimpleNamespace(model_type="llama"), + sparse_attention_config=None, + ) + llm_args = SimpleNamespace( + encode_only=False, + mm_encoder_only=False, + get_runtime_sizes=lambda: (1, 16, 64, max_batch_size), + encoder_max_batch_size=None, + encoder_max_num_tokens=None, + enable_in_graph_sampling=False, + enable_return_routed_experts=False, + multimodal_config=SimpleNamespace(video_pruning_rate=None), + checkpoint_format="HF", + trust_remote_code=False, + disable_overlap_scheduler=False, + kv_cache_config=KvCacheConfig(), + sparse_attention_config=None, + cache_transceiver_config=cache_transceiver_config, + ) + engine = PyTorchModelEngine.__new__(PyTorchModelEngine) + with pytest.raises(RuntimeError, match="seat pool sized"): + PyTorchModelEngine.__init__( + engine, + model_path="dummy", + mapping=Mapping(), + model=model, + llm_args=llm_args, + checkpoint_loader=Mock(), + ) + return engine.max_num_seq_slots, engine.disagg_gen_prefetch_requests + + +_NIXL = SimpleNamespace(backend="NIXL") + + +@pytest.mark.parametrize("max_batch_size", [1, 2]) +def test_engine_seat_pool_is_unchanged_by_default(monkeypatch, prefetch_env, max_batch_size): + prefetch_env(None) + for transceiver in (None, _NIXL): + assert _engine_seat_pool( + monkeypatch, cache_transceiver_config=transceiver, max_batch_size=max_batch_size + ) == (max_batch_size, 0) + + +@pytest.mark.parametrize("prefetch", [1, 3]) +@pytest.mark.parametrize("max_batch_size", [1, 2]) +def test_engine_seat_pool_adds_the_prefetch_seats_for_disagg( + monkeypatch, prefetch_env, max_batch_size, prefetch +): + prefetch_env(str(prefetch)) + assert _engine_seat_pool( + monkeypatch, cache_transceiver_config=_NIXL, max_batch_size=max_batch_size + ) == (max_batch_size + prefetch, prefetch) + + +def test_engine_seat_pool_ignores_prefetch_without_a_transceiver(monkeypatch, prefetch_env): + prefetch_env("1") + assert _engine_seat_pool(monkeypatch, cache_transceiver_config=None) == (2, 0) + # A transceiver config without a backend is not disaggregated either. + assert _engine_seat_pool( + monkeypatch, cache_transceiver_config=SimpleNamespace(backend=None) + ) == (2, 0) + + +def test_engine_seat_pool_ignores_prefetch_with_prioritize_first_token(monkeypatch, prefetch_env): + prefetch_env("1") + monkeypatch.setenv(PRIORITIZE_FIRST_TOKEN_ENV, "1") + assert _engine_seat_pool(monkeypatch, cache_transceiver_config=_NIXL) == (2, 0) + + +def _seq_req(request_id, *, transmission_complete=False): + """The attributes ``SeqSlotManager.prepare_resources`` reads and writes.""" + return SimpleNamespace( + request_id=request_id, + seq_slot=None, + py_seq_slot=None, + return_perf_metrics=False, + is_disagg_generation_init_state=False, + is_disagg_generation_transmission_complete=transmission_complete, + ) + + +def _scheduled(*requests): + return SimpleNamespace(all_requests=lambda: list(requests)) + + +def _overlap_retire_turnover(pool, max_batch_size): + """Seat ``max_batch_size`` decodes, then schedule a prefetched request. + + Under the overlap scheduler the request that ended on the previous step is + GENERATION_TO_COMPLETE: it is not scheduled but still owns its sequence + slot until the next iteration frees it. With prefetch, the request whose KV + already landed is scheduled in that same iteration. + """ + manager = SeqSlotManager(max_num_sequences=pool) + running = [_seq_req(i) for i in range(max_batch_size)] + manager.prepare_resources(_scheduled(*running)) + prefetched = _seq_req(100, transmission_complete=True) + manager.prepare_resources(_scheduled(*running[1:], prefetched)) + return manager, running, prefetched + + +@pytest.mark.parametrize("max_batch_size", [1, 2, 8]) +def test_a_seat_pool_of_max_batch_size_cannot_seat_the_prefetched_request(max_batch_size): + with pytest.raises(NoFreeSlotsError): + _overlap_retire_turnover(max_batch_size, max_batch_size) + + +@pytest.mark.parametrize("prefetch", [1, 2]) +@pytest.mark.parametrize("max_batch_size", [1, 2, 8]) +def test_the_widened_seat_pool_seats_the_prefetched_request(max_batch_size, prefetch): + pool = max_batch_size + prefetch + manager, running, prefetched = _overlap_retire_turnover(pool, max_batch_size) + assert 0 <= prefetched.py_seq_slot < pool + assert prefetched.py_seq_slot not in {r.py_seq_slot for r in running} + # The retiring request's slot is returned and reused on later turnovers. + manager.free_resources(running[0]) + assert len(manager.slot_manager.free_slots) == pool - max_batch_size + + +def _executor(max_num_active_requests=1, prefetch=0, attention_dp=False): + ex = object.__new__(PyExecutor) + ex.enable_attention_dp = attention_dp + ex.max_num_active_requests = max_num_active_requests + ex.disagg_gen_prefetch_requests = prefetch + ex.is_benchmark_disagg = False + ex._benchmark_fill_phase_active = False + ex.dist = Mock() + ex.dist.tp_size = 4 + return ex + + +@pytest.mark.parametrize("mbs", [1, 2, 3]) +def test_admission_capacity_gate(mbs): + assert _executor(mbs, prefetch=0)._get_request_admission_capacity() == mbs + assert _executor(mbs, prefetch=1)._get_request_admission_capacity() == mbs + 1 + assert _executor(mbs, prefetch=2)._get_request_admission_capacity() == mbs + 2 + + +def test_admission_capacity_unchanged_without_the_attribute(): + ex = _executor(2) + del ex.disagg_gen_prefetch_requests + assert ex._get_request_admission_capacity() == 2 + + +def test_admission_capacity_ignores_prefetch_under_attention_dp(): + ex = _executor(2, prefetch=1, attention_dp=True) + assert ex._get_request_admission_capacity() == 4 * 2 + + +def _item(req_id, request_type, child_req_ids=None): + return RequestQueueItem( + req_id, + SimpleNamespace(request_type=request_type, priority=DEFAULT_REQUEST_PRIORITY), + child_req_ids=child_req_ids, + ) + + +GEN = RequestType.REQUEST_TYPE_GENERATION_ONLY +CTX = RequestType.REQUEST_TYPE_CONTEXT_ONLY + + +def _pop(ex, items, num_active): + queue = FCFSWaitingQueue() + queue.add_requests(items) + popped = ex._pop_from_waiting_queue(queue, num_active) + return [i.id for i in popped], [i.id for i in queue] + + +def test_gate_off_pops_up_to_max_batch_size(): + ex = _executor(1, prefetch=0) + assert _pop(ex, [_item(1, GEN), _item(2, GEN)], 1) == ([], [1, 2]) + assert _pop(ex, [_item(1, GEN), _item(2, GEN)], 0) == ([1], [2]) + + +def test_prefetch_admits_one_more_generation_request(): + ex = _executor(1, prefetch=1) + # One request decoding: the next one is admitted to receive its KV. + assert _pop(ex, [_item(1, GEN), _item(2, GEN)], 1) == ([1], [2]) + # Idle: one regular slot and one prefetch slot. + assert _pop(ex, [_item(1, GEN), _item(2, GEN), _item(3, GEN)], 0) == ([1, 2], [3]) + # Full including the prefetch slot. + assert _pop(ex, [_item(1, GEN)], 2) == ([], [1]) + + +def test_prefetch_slots_take_only_generation_only_requests(): + ex = _executor(1, prefetch=2) + # A context request cannot use a prefetch slot; it and everything after + # it go back to the front of the queue in order. + assert _pop(ex, [_item(1, CTX), _item(2, GEN)], 1) == ([], [1, 2]) + assert _pop(ex, [_item(1, CTX), _item(2, CTX), _item(3, GEN)], 0) == ([1], [2, 3]) + assert _pop(ex, [_item(1, GEN), _item(2, CTX), _item(3, GEN)], 0) == ([1], [2, 3]) + assert _pop(ex, [_item(1, GEN), _item(2, GEN), _item(3, CTX)], 0) == ([1, 2], [3]) + # Requests with children never use a prefetch slot. + assert _pop(ex, [_item(1, GEN, child_req_ids=[9])], 1) == ([], [1]) + + +class _FakeDFlashWorker: + """The DFlash worker state ``DFlashSpecMetadata.prepare`` reads and writes. + + ``_assign_slot`` is the real method; ``_write_ctx_len`` keeps the host + mirror (``ctx_len``) that the real one keeps next to the device tensor. + """ + + def __init__(self, max_batch_size): + self._ctx_buf_inited = True + self._req_to_slot = {} + self._req_ctx_pos = {} + self._free_slots = deque(range(max_batch_size)) + self._dummy_slot = max_batch_size + self._graph_dummy_id_floor = 1 << 62 + self._batch_to_slot = torch.zeros(max_batch_size, dtype=torch.long) + self.ctx_len = [0] * (max_batch_size + 1) + self._assign_slot = MethodType(DFlashWorker._assign_slot, self) + + def _write_ctx_len(self, updates, base_updates=None): + for slot, value in updates.items(): + self.ctx_len[slot] = value + + +def _drafter_slots(monkeypatch, max_batch_size, schedule, stored_context=500): + """Run ``DFlashSpecMetadata.prepare`` over a schedule of generation batches. + + After each step, every request in the batch is given ``stored_context`` + tokens of drafter context, as the real step does. Returns, per step, each + request's slot and the context length it found when the step started. + """ + monkeypatch.setattr(dflash_module, "prefer_pinned", lambda: False) + worker = _FakeDFlashWorker(max_batch_size) + metadata = object.__new__(DFlashSpecMetadata) + metadata.batch_indices_cuda = torch.zeros(max_batch_size, dtype=torch.int) + metadata._dflash_worker = worker + seen = [] + for batch in schedule: + metadata.request_ids = list(batch) + metadata.num_generations = len(batch) + metadata.prepare() + slots = worker._batch_to_slot[: len(batch)].tolist() + seen.append({rid: (slot, worker.ctx_len[slot]) for rid, slot in zip(batch, slots)}) + for slot in slots: + worker.ctx_len[slot] = stored_context + return seen + + +def test_prefetch_turnover_gives_the_drafter_the_slot_state_of_a_plain_turnover(monkeypatch): + """A prefetched request starts from the same drafter slot state as without prefetch. + + Request 11 ends while request 12 keeps decoding. Without prefetch the next request + (13) is scheduled after at least one more step (here, one step with + request 12 alone); with prefetch it is scheduled on the very next step. + DFlash frees the slots of requests that left the batch before it assigns + slots to new ones, so request 13 gets the same slot, reset to zero + context, either way, and request 12 keeps its own slot and context. + (Request id 0 is the attention-DP dummy, so real ids start above it.) + """ + today = _drafter_slots(monkeypatch, 2, [[11, 12], [12], [12, 13]]) + prefetch = _drafter_slots(monkeypatch, 2, [[11, 12], [12, 13]]) + assert today[-1] == prefetch[-1] + slot_13, context_13 = prefetch[-1][13] + assert context_13 == 0 + assert slot_13 == prefetch[0][11][0] # the retired request's slot + assert slot_13 != 2 # not the shared dummy slot + assert prefetch[-1][12] == (prefetch[0][12][0], 500) + + +def test_prefetch_turnover_at_max_batch_size_one_reuses_the_slot_from_zero(monkeypatch): + # With one slot the sequence of drafter batches is [11] then [12] with and + # without prefetch; request 12 takes request 11's slot at zero context. + seen = _drafter_slots(monkeypatch, 1, [[11], [12]]) + assert seen[1][12] == (seen[0][11][0], 0) diff --git a/tests/unittest/disaggregated/test_router_hold.py b/tests/unittest/disaggregated/test_router_hold.py new file mode 100644 index 000000000000..6fcd650d4710 --- /dev/null +++ b/tests/unittest/disaggregated/test_router_hold.py @@ -0,0 +1,425 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Held placement in LoadBalancingRouter (``max_inflight_per_server``). + +Covers the in-process router, the coordinator-owned router reached through +``DisaggCoordinatorService.select`` / ``finish``, and the fleet-worker side +``CoordinatorDelegatingRouter``. +""" + +import asyncio +from unittest import mock + +import pytest + +from tensorrt_llm.llmapi.disagg_utils import ( + CtxGenServerConfig, + DisaggServerConfig, + RouterConfig, + ServerRole, + extract_router_config, +) +from tensorrt_llm.serve.disagg_coordinator import DisaggCoordinatorService +from tensorrt_llm.serve.openai_protocol import CompletionRequest, DisaggregatedParams +from tensorrt_llm.serve.router import ( + PLACEMENT_HOLD_TIMEOUT_DEFAULT_S, + CoordinatorDelegatingRouter, + HeldPlacementCancelled, + LoadBalancingRouter, + create_router, +) + +pytestmark = pytest.mark.cpu_only + + +def _request(i: int = 0) -> CompletionRequest: + return CompletionRequest(model="m", prompt=f"prompt {i}") + + +def _router(servers, cap=None, **kwargs) -> LoadBalancingRouter: + return LoadBalancingRouter( + server_role=ServerRole.GENERATION, + servers=list(servers), + max_inflight_per_server=cap, + **kwargs, + ) + + +def _inflight(router: LoadBalancingRouter) -> dict: + return {server: state._num_active_requests for server, state in router._server_state.items()} + + +async def _settle(): + for _ in range(5): + await asyncio.sleep(0) + + +# ---- gate off ------------------------------------------------------------------ + + +@pytest.mark.asyncio +async def test_default_never_holds(): + router = _router(["g0", "g1"]) + assert router.placement_hold_timeout_s is None + requests = [_request(i) for i in range(5)] + placed = [(await router.get_next_server(r))[0] for r in requests] + # JSQ with round-robin ties: alternates and never waits. + assert placed == ["g0", "g1", "g0", "g1", "g0"] + assert _inflight(router) == {"g0": 3, "g1": 2} + for r in requests: + await router.finish_request(r) + assert _inflight(router) == {"g0": 0, "g1": 0} + + +def test_router_config_reaches_the_router(): + server_cfg = { + "router": {"type": "load_balancing", "max_inflight_per_server": 2, "hold_timeout_s": 30}, + "max_batch_size": 1, + } + config = extract_router_config(server_cfg) + config.server_role = ServerRole.GENERATION + router = create_router(config, ["g0", "g1"]) + assert isinstance(router, LoadBalancingRouter) + assert router._max_inflight_per_server == 2 + assert router.placement_hold_timeout_s == 30.0 + + plain = create_router( + RouterConfig(type="load_balancing", server_role=ServerRole.GENERATION, args={}), ["g0"] + ) + assert plain.placement_hold_timeout_s is None + assert PLACEMENT_HOLD_TIMEOUT_DEFAULT_S > 0 + + +@pytest.mark.parametrize("cap", [0, -1, 1.5, True, "2"]) +def test_invalid_cap_is_rejected(cap): + with pytest.raises(ValueError): + _router(["g0"], cap=cap) + + +def test_invalid_hold_timeout_is_rejected(): + with pytest.raises(ValueError): + _router(["g0"], cap=1, hold_timeout_s=0) + + +# ---- in-process hold ----------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_hold_and_release_are_fifo(): + router = _router(["g0", "g1"], cap=1) + r = [_request(i) for i in range(5)] + s0, _ = await router.get_next_server(r[0]) + s1, _ = await router.get_next_server(r[1]) + assert {s0, s1} == {"g0", "g1"} + + tasks = [] + for i in (2, 3, 4): + tasks.append(asyncio.create_task(router.get_next_server(r[i]))) + await _settle() + assert router.num_held_requests == 3 + assert not any(t.done() for t in tasks) + assert _inflight(router) == {"g0": 1, "g1": 1} + + await router.finish_request(r[1]) + await _settle() + assert tasks[0].done() and not tasks[1].done() + assert tasks[0].result()[0] == s1 + + await router.finish_request(r[0]) + await _settle() + assert tasks[1].done() and not tasks[2].done() + assert tasks[1].result()[0] == s0 + assert router.num_held_requests == 1 + + await router.finish_request(r[2]) + await _settle() + assert tasks[2].result()[0] == s1 + assert _inflight(router) == {"g0": 1, "g1": 1} + for i in (3, 4): + await router.finish_request(r[i]) + assert _inflight(router) == {"g0": 0, "g1": 0} + assert router._req_routing_table == {} + + +@pytest.mark.asyncio +async def test_prefers_idle_then_youngest_running(): + router = _router(["g0", "g1", "g2"], cap=2) + r = [_request(i) for i in range(4)] + first = [] + for x in r[:3]: + first.append((await router.get_next_server(x))[0]) + await asyncio.sleep(0.002) + # Idle servers first: one request per server before any gets a second. + assert sorted(first) == ["g0", "g1", "g2"] + # All have one in flight: the server whose request started last wins. + s3, _ = await router.get_next_server(r[3]) + assert s3 == first[2] + + # A new request never waits while a server is below the cap. + await router.finish_request(r[0]) + s4, _ = await router.get_next_server(_request(4)) + assert s4 == first[0] + + +@pytest.mark.asyncio +async def test_new_arrival_waits_behind_held_requests(): + router = _router(["g0"], cap=1) + r = [_request(i) for i in range(3)] + await router.get_next_server(r[0]) + t1 = asyncio.create_task(router.get_next_server(r[1])) + await _settle() + await router.finish_request(r[0]) + # The slot freed by r0 goes to r1, which was waiting, not to r2. + t2 = asyncio.create_task(router.get_next_server(r[2])) + await _settle() + assert t1.done() and not t2.done() + await router.finish_request(r[1]) + await _settle() + assert t2.done() + await router.finish_request(r[2]) + assert _inflight(router) == {"g0": 0} + + +@pytest.mark.asyncio +async def test_cancel_while_held_releases_cleanly(): + router = _router(["g0"], cap=1) + r = [_request(i) for i in range(3)] + await router.get_next_server(r[0]) + t1 = asyncio.create_task(router.get_next_server(r[1])) + t2 = asyncio.create_task(router.get_next_server(r[2])) + await _settle() + assert router.num_held_requests == 2 + + t1.cancel() + with pytest.raises(asyncio.CancelledError): + await t1 + assert router.num_held_requests == 1 + + await router.finish_request(r[0]) + await _settle() + assert t2.result()[0] == "g0" + assert _inflight(router) == {"g0": 1} + await router.finish_request(r[2]) + assert _inflight(router) == {"g0": 0} + assert router._req_routing_table == {} + assert router._hold_started == {} + + +@pytest.mark.asyncio +async def test_cancel_after_assignment_gives_the_slot_back(): + router = _router(["g0"], cap=1) + r = [_request(i) for i in range(3)] + await router.get_next_server(r[0]) + t1 = asyncio.create_task(router.get_next_server(r[1])) + t2 = asyncio.create_task(router.get_next_server(r[2])) + await _settle() + # r1 is assigned g0 and cancelled in the same loop turn. + await router.finish_request(r[0]) + t1.cancel() + with pytest.raises(asyncio.CancelledError): + await t1 + await _settle() + # The slot r1 never used went to r2. + assert t2.result()[0] == "g0" + assert _inflight(router) == {"g0": 1} + await router.finish_request(r[2]) + assert _inflight(router) == {"g0": 0} + + +@pytest.mark.asyncio +async def test_hold_timeout_falls_back_to_least_loaded(): + router = _router(["g0", "g1"], cap=1, hold_timeout_s=0.05) + r = [_request(i) for i in range(3)] + await router.get_next_server(r[0]) + await router.get_next_server(r[1]) + server, _ = await asyncio.wait_for(router.get_next_server(r[2]), 5) + assert server in ("g0", "g1") + assert router._hold_fallbacks == 1 + assert router.num_held_requests == 0 + assert sum(_inflight(router).values()) == 3 + for x in r: + await router.finish_request(x) + assert _inflight(router) == {"g0": 0, "g1": 0} + + +@pytest.mark.asyncio +async def test_new_server_releases_held_requests(): + router = _router(["g0"], cap=1) + r = [_request(i) for i in range(2)] + await router.get_next_server(r[0]) + t1 = asyncio.create_task(router.get_next_server(r[1])) + await _settle() + old = router._servers.copy() + router._servers = old + ["g1"] + router._on_servers_updated(old, router._servers) + await _settle() + assert t1.result()[0] == "g1" + + +# ---- coordinator-owned router (by request ID) ------------------------------------ + + +@pytest.mark.asyncio +async def test_by_id_finish_cancels_a_held_placement(): + router = _router(["g0"], cap=1) + await router.get_next_server_by_key({}, req_id=1) + held = asyncio.create_task(router.get_next_server_by_key({}, req_id=2)) + await _settle() + await router.finish_request_by_id(2, False) + with pytest.raises(HeldPlacementCancelled): + await held + assert router.num_held_requests == 0 + await router.finish_request_by_id(1) + assert _inflight(router) == {"g0": 0} + + +@pytest.mark.asyncio +async def test_by_id_release_before_placement_refuses_the_late_select(): + router = _router(["g0"], cap=1) + await router.finish_request_by_id(7, False) + with pytest.raises(HeldPlacementCancelled): + await router.get_next_server_by_key({}, req_id=7) + assert _inflight(router) == {"g0": 0} + # The marker is consumed: a later placement for the ID goes through. + server, _, rid = await router.get_next_server_by_key({}, req_id=7) + assert (server, rid) == ("g0", 7) + + +@pytest.mark.asyncio +async def test_by_id_repeated_placement_does_not_leak_load(): + router = _router(["g0", "g1"], cap=2, use_tokens=True) + await router.get_next_server_by_key({"num_tokens": 10}, req_id=5) + await router.get_next_server_by_key({"num_tokens": 10}, req_id=5) + assert sum(_inflight(router).values()) == 1 + await router.finish_request_by_id(5) + assert _inflight(router) == {"g0": 0, "g1": 0} + assert all(state._num_active_tokens == 0 for state in router._server_state.values()) + + +def _coordinator(gen_urls, gen_router_args): + config = DisaggServerConfig( + server_configs=[ + CtxGenServerConfig(type="gen", hostname=u.split(":")[0], port=int(u.split(":")[1])) + for u in gen_urls + ], + ctx_router_config=RouterConfig(type="round_robin", server_role=ServerRole.CONTEXT), + gen_router_config=RouterConfig( + type="load_balancing", server_role=ServerRole.GENERATION, args=dict(gen_router_args) + ), + ) + return DisaggCoordinatorService(config, lambda *a, **k: None, reservation_timeout_secs=60) + + +def _same_stripe_id(coordinator, req_id): + stripe = coordinator._reservation_lock(("generation", req_id)) + for other in range(req_id + 1, req_id + 100000): + if coordinator._reservation_lock(("generation", other)) is stripe: + return other + raise AssertionError("no colliding request ID found") + + +@pytest.mark.asyncio +async def test_coordinator_hold_releases_on_finish_sharing_the_lock_stripe(): + coordinator = _coordinator(["gen0:8000"], {"max_inflight_per_server": 1}) + first = 11 + held_id = _same_stripe_id(coordinator, first) + await coordinator.select("generation", {}, first, None) + held = asyncio.create_task(coordinator.select("generation", {}, held_id, None)) + await _settle() + assert coordinator.gen_router.num_held_requests == 1 + # The finish takes the same reservation-lock stripe as the held select. + await asyncio.wait_for(coordinator.finish("generation", first), 5) + server, _, rid = await asyncio.wait_for(held, 5) + assert (server, rid) == ("gen0:8000", held_id) + assert ("generation", held_id) in coordinator._reservation_tasks + await coordinator.finish("generation", held_id) + assert coordinator._reservation_tasks == {} + assert _inflight(coordinator.gen_router) == {"gen0:8000": 0} + + +@pytest.mark.asyncio +async def test_coordinator_finish_of_a_held_select_fails_it(): + coordinator = _coordinator(["gen0:8000"], {"max_inflight_per_server": 1}) + await coordinator.select("generation", {}, 1, None) + held = asyncio.create_task(coordinator.select("generation", {}, 2, None)) + await _settle() + await coordinator.finish("generation", 2, False) + with pytest.raises(ValueError): + await held + assert ("generation", 2) not in coordinator._reservation_tasks + await coordinator.finish("generation", 1) + assert _inflight(coordinator.gen_router) == {"gen0:8000": 0} + + +@pytest.mark.asyncio +async def test_coordinator_concurrent_duplicate_select_with_hold(): + coordinator = _coordinator(["gen0:8000", "gen1:8000"], {"max_inflight_per_server": 1}) + await asyncio.gather( + coordinator.select("generation", {}, 123, None), + coordinator.select("generation", {}, 123, None), + ) + assert sum(_inflight(coordinator.gen_router).values()) == 1 + await coordinator.finish("generation", 123) + assert sum(_inflight(coordinator.gen_router).values()) == 0 + assert coordinator._reservation_tasks == {} + + +# ---- fleet-worker side ----------------------------------------------------------- + + +def _delegating(local): + return CoordinatorDelegatingRouter( + "http://coordinator", local, "generation", request_timeout_s=5.0 + ) + + +def _gen_request(): + request = _request() + request.disaggregated_params = DisaggregatedParams( + request_type="generation_only", ctx_request_id=42, disagg_request_id=42 + ) + return request + + +@pytest.mark.asyncio +async def test_worker_select_timeout_default_is_unchanged(): + router = _delegating(_router(["g0"])) + router._post_select = mock.AsyncMock(return_value={"server": "g0", "info": {}}) + server, _ = await router.get_next_server(_gen_request()) + assert server == "g0" + assert router._post_select.await_args.kwargs == {} + + +@pytest.mark.asyncio +async def test_worker_select_waits_for_the_hold_and_releases_on_cancel(): + router = _delegating(_router(["g0"], cap=1, hold_timeout_s=120)) + router._post_select = mock.AsyncMock(return_value={"server": "g0", "info": {}}) + await router.get_next_server(_gen_request()) + assert router._post_select.await_args.kwargs == {"timeout": 125.0} + + started = asyncio.Event() + + async def _hang(payload, timeout=None): + started.set() + await asyncio.sleep(3600) + + router._post_select = _hang + router._ensure_finish_workers = mock.Mock() + task = asyncio.create_task(router.get_next_server(_gen_request())) + await started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert router._finish_queue.get_nowait() == (42, False) From 357c1be39a8fc0e5a66a90a40ab7637474222b7d Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:10:56 +0000 Subject: [PATCH 04/15] [None][feat] PyExecutor: model-declared runtime hook to run the overlap loop on the execution stream RuntimeHooks gains step_loop_on_execution_stream (False by default). When a model declares it, the overlap loop's thread makes the execution stream current, so the sampler and the next step's input preparation queue behind the forward on one stream instead of crossing streams twice per step. Models that do not declare it are unchanged. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 21 ++++++++++++++++++- .../_torch/pyexecutor/runtime_hooks.py | 8 +++++++ .../_torch/executor/test_runtime_hooks.py | 3 ++- 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 929b93ac677e..37a09079a8bd 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -113,6 +113,7 @@ get_from_waiting_queue, merge_requests) from .resource_manager import (NoFreeSlotsError, ResourceManager, ResourceManagerType, request_context) +from .runtime_hooks import runtime_hooks_of from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState, SampleStateTensors) from .scheduler import (RequestScheduler, ScheduledRequests, @@ -5441,6 +5442,23 @@ def _wait_for_model_engine_input_copy(self): if wait_for_input_copy is not None: wait_for_input_copy() + def _step_loop_stream(self): + """The stream the overlap loop's thread makes current. + + ``nullcontext()`` (the thread's own current stream) unless the model's + runtime hooks declare ``step_loop_on_execution_stream``; then the + execution stream, so the sampler and the next step's input preparation + queue behind the forward on one stream (see ``RuntimeHooks``). + """ + model = getattr(self.model_engine, "model", None) + if model is not None and runtime_hooks_of( + model).step_loop_on_execution_stream: + logger.info( + "[PyExecutor] overlap loop runs on the execution stream " + f"{self.execution_stream} (runtime hook).") + return torch.cuda.stream(self.execution_stream) + return nullcontext() + def _executor_loop_overlap(self): torch.cuda.set_device(self.device_id) # ensure the context is created, otherwise, some MPI calls will fail. @@ -5448,7 +5466,8 @@ def _executor_loop_overlap(self): self._kv_defer_close_armed = ( self._is_kv_manager_v2 and self.kv_cache_manager.arm_deferred_close(True)) - with self._profiler() as profile_step, self.hang_detector: + with self._step_loop_stream(), self._profiler( + ) as profile_step, self.hang_detector: iter_start_time = time.time() iter_stats = None target_inputs = None diff --git a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py index 5b72fefa9352..b00ed0e6ec58 100644 --- a/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py +++ b/tensorrt_llm/_torch/pyexecutor/runtime_hooks.py @@ -143,6 +143,14 @@ class RuntimeHooks: step_input_stager: Optional[StepInputStagerSource] = None # The speculative sampler's store update and host copy in one launch (see SamplerStoreWriter). sampler_store_writer: Optional[SamplerStoreWriter] = None + # Run the overlap executor loop with the execution stream as its thread's current stream. + # By default the forward runs on the execution stream and the sampler and the next + # step's input preparation on the thread's own stream, ordered by two cross-stream + # waits per step (the sampler waits for the forward, the next forward waits for the + # sampler); on one stream the same work runs in the same host order with no + # cross-stream wait between two forwards. Every dependency the waits expressed holds + # on one stream, so only overlap between the two streams is given up. + step_loop_on_execution_stream: bool = False DEFAULT_RUNTIME_HOOKS = RuntimeHooks() diff --git a/tests/unittest/_torch/executor/test_runtime_hooks.py b/tests/unittest/_torch/executor/test_runtime_hooks.py index 842739e4fce4..8112c20cacc8 100644 --- a/tests/unittest/_torch/executor/test_runtime_hooks.py +++ b/tests/unittest/_torch/executor/test_runtime_hooks.py @@ -62,7 +62,8 @@ def test_a_model_without_hooks_gets_the_all_off_default(model): def test_every_field_defaults_to_off(): for field in dataclasses.fields(RuntimeHooks): - assert getattr(DEFAULT_RUNTIME_HOOKS, field.name) is None, field.name + # Object hooks are off as None, flags as False. + assert getattr(DEFAULT_RUNTIME_HOOKS, field.name) in (None, False), field.name def test_hooks_declared_as_none_mean_the_default(): From 3aa02730b5b6d0c3741059940c6fac4010296868 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:10:57 +0000 Subject: [PATCH 05/15] [None][feat] Muse Spark Open: opt-in NVFP4 experts for the quantization-excluded MoE layers The checkpoint keeps the experts of its quantization-excluded MoE layers (0 and 61) in BF16. With MUSESPARK_OPEN_L061_EXPERTS_NVFP4=1 (off by default) those experts are built on the global NVFP4 leaf and quantized from BF16 at load. MUSESPARK_OPEN_L061_EXPERTS_NVFP4_LAYERS selects a subset of the layers, and MUSESPARK_OPEN_L061_EXPERTS_NVFP4_FC2_AMAX supplies the FC2 input amax per layer. This is a precision change and is never enabled implicitly. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../hf/musespark_open_weight_mapper.py | 176 +++++++- .../_torch/models/modeling_musespark_open.py | 74 +++- .../test_musespark_open_l061_nvfp4.py | 395 ++++++++++++++++++ 3 files changed, 633 insertions(+), 12 deletions(-) create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_l061_nvfp4.py diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/musespark_open_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/musespark_open_weight_mapper.py index a0b53e24de23..ef7a5185f1f5 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/musespark_open_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/musespark_open_weight_mapper.py @@ -31,6 +31,7 @@ expert range and the rank's attention shard are ever materialized. """ +import os import re from typing import Any, Dict, List, Optional, Tuple @@ -38,7 +39,11 @@ from tensorrt_llm.logger import logger -from ...modeling_musespark_open import PERCEPTION_KEY_PREFIXES, head_interleaved_q_gate +from ...modeling_musespark_open import ( + PERCEPTION_KEY_PREFIXES, + excluded_layer_experts_nvfp4_enabled, + head_interleaved_q_gate, +) from .musespark_open_format import OPEN_FORMAT_SENTINEL_KEY, build_open_format_rename _LAYER_KEY = re.compile(r"^model\.layers\.(\d+)\.") @@ -411,6 +416,126 @@ def build_expert_weights( return entries +#: E2M1 max (6) times E4M3 max (448). An NVFP4 global scale maps a tensor's +#: amax onto it; the checkpoint stores the inverse (``amax / 2688``) as +#: ``weight_scale_2`` and ``input_scale``. +NVFP4_GLOBAL_RANGE = 6.0 * 448.0 + +#: Elements per NVFP4 block scale (the checkpoint's ``group_size``). +NVFP4_BLOCK = 16 + + +def nvfp4_quantize_weight(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize one ``[rows, cols]`` weight to the checkpoint's NVFP4 layout. + + Returns ``(packed, block_scale, weight_scale_2)``: ``packed`` is ``uint8 + [rows, cols / 2]`` (two E2M1 values per byte, the lower column in the low + nibble), ``block_scale`` is ``float8_e4m3fn [rows, cols / 16]`` in row-major + (unswizzled) order, and ``weight_scale_2`` is the FP32 0-d per-tensor scale + ``amax / 2688``. Dequantization is ``e2m1 * block_scale * weight_scale_2``, + the same as for the checkpoint's quantized layers; the NVFP4 MoE method + swizzles and shuffles these at load exactly as it does theirs. + """ + rows, cols = weight.shape + source = weight.to(device="cuda", dtype=torch.bfloat16).contiguous() + amax = source.abs().amax().float().clamp_min(torch.finfo(torch.float32).tiny) + global_scale = (NVFP4_GLOBAL_RANGE / amax).reshape(1) + packed, block_scale = torch.ops.trtllm.fp4_quantize( + source, global_scale, NVFP4_BLOCK, False, False + ) + block_scale = block_scale.view(torch.uint8).reshape(rows, cols // NVFP4_BLOCK) + return ( + packed.reshape(rows, cols // 2), + block_scale.view(torch.float8_e4m3fn), + (amax / NVFP4_GLOBAL_RANGE).cpu().reshape([]), + ) + + +def excluded_layer_fc2_amax(layer_idx: int) -> float: + """The calibrated FC2-input amax of a quant-excluded layer. + + Read from MUSESPARK_OPEN_L061_EXPERTS_NVFP4_FC2_AMAX, a comma-separated + list of ``layer=amax`` pairs. The FC2 input (the SwiGLU product) has no + bound that follows from the layer's weights alone, so its static NVFP4 + global scale needs a measured value, as the checkpoint's quantized layers + have. + + Raises: + ValueError: if the variable does not name ``layer_idx``. + """ + spec = os.environ.get("MUSESPARK_OPEN_L061_EXPERTS_NVFP4_FC2_AMAX", "") + for item in spec.split(","): + if "=" not in item: + continue + layer, value = item.split("=", 1) + if int(layer) == layer_idx: + amax = float(value) + if not amax > 0.0: + raise ValueError(f"MuseSparkOpen: FC2 amax of layer {layer_idx} must be > 0") + return amax + raise ValueError( + f"MuseSparkOpen: layer {layer_idx} experts are quantized at load, but " + "MUSESPARK_OPEN_L061_EXPERTS_NVFP4_FC2_AMAX names no FC2 amax for it" + ) + + +def build_nvfp4_expert_weights_from_bf16( + weights, + prefix: str, + expert_start: int, + expert_end: int, + num_experts: int, + intermediate_size: int, + fc1_amax: float, + fc2_amax: float, + accounting: WeightAccounting, +) -> Dict[str, Any]: + """The entries :func:`build_expert_weights` yields for a quantized layer, from BF16 experts. + + Quantizes this rank's ``[expert_start, expert_end)`` experts of a layer the + checkpoint stores in BF16 into the checkpoint's NVFP4 entry layout: per + expert and projection an E2M1 ``weight``, its E4M3 ``weight_scale`` and its + FP32 ``weight_scale_2``, with gate (``w1``) and up (``w3``) scaled + separately as the checkpoint's per-half ``w13_weight_scale_2`` does. Every + expert gets the same activation scales, ``fc1_amax / 2688`` for both FC1 + halves and ``fc2_amax / 2688`` for FC2, so the layer-level maxima the MoE + method takes are those two values on every rank. + """ + entries: Dict[str, Any] = {} + w13 = _slice_experts(weights, f"{prefix}.w13_weight", expert_start, expert_end, accounting) + w2 = _slice_experts(weights, f"{prefix}.w2_weight", expert_start, expert_end, accounting) + for local_idx, expert_id in enumerate(range(expert_start, expert_end)): + for name, weight in ( + ("w1", w13[local_idx, :intermediate_size]), + ("w3", w13[local_idx, intermediate_size:]), + ("w2", w2[local_idx]), + ): + packed, block_scale, weight_scale_2 = nvfp4_quantize_weight(weight) + entries[f"{expert_id}.{name}.weight"] = packed + entries[f"{expert_id}.{name}.weight_scale"] = block_scale + entries[f"{expert_id}.{name}.weight_scale_2"] = weight_scale_2 + fc1_input_scale = torch.tensor(fc1_amax / NVFP4_GLOBAL_RANGE, dtype=torch.float32) + fc2_input_scale = torch.tensor(fc2_amax / NVFP4_GLOBAL_RANGE, dtype=torch.float32) + for expert_id in range(num_experts): + entries[f"{expert_id}.w1.input_scale"] = fc1_input_scale + entries[f"{expert_id}.w3.input_scale"] = fc1_input_scale + entries[f"{expert_id}.w2.input_scale"] = fc2_input_scale + return entries + + +def pre_expert_norm_amax(norm_weight: torch.Tensor, hidden_size: int) -> float: + """An upper bound on ``|RMSNorm(x)|`` for the Gemma-style pre-expert norm. + + The norm computes ``x / rms(x) * gain`` with ``gain = bf16(weight + 1)``. + Since ``|x_i| <= sqrt(n) * rms(x)`` for any ``n``-element ``x`` (equality + only for a one-hot ``x``; the epsilon only lowers it), every output element + is at most ``sqrt(n) * max|gain|``. Used as the FC1 activation amax of the + layers quantized at load: no input can saturate the NVFP4 block scales. + """ + gain = (norm_weight.float() + 1.0).to(torch.bfloat16).float() + return float(hidden_size**0.5 * gain.abs().max().item()) + + def load_layer( model, layer, @@ -488,16 +613,45 @@ def load_layer( experts_prefix = f"{ffn_prefix}.experts" is_quantized = f"{experts_prefix}.w13_weight_scale" in weights expert_start, expert_end = _local_expert_range(moe.experts, config.num_experts) - expert_weights = build_expert_weights( - weights, - experts_prefix, - expert_start, - expert_end, - config.num_experts, - config.intermediate_size, - is_quantized, - accounting, - ) + experts_backend = getattr(moe.experts, "backend", moe.experts) + if not is_quantized and getattr(experts_backend, "has_nvfp4", False): + if not excluded_layer_experts_nvfp4_enabled(): + raise ValueError( + f"MuseSparkOpen: layer {layer_idx} experts are BF16 in the checkpoint " + "but were built NVFP4" + ) + fc1_amax = pre_expert_norm_amax(moe.pre_expert_norm.weight, config.expert_hidden_size) + fc2_amax = excluded_layer_fc2_amax(layer_idx) + logger.info( + f"MuseSparkOpen: quantizing layer {layer_idx} experts " + f"[{expert_start}, {expert_end}) to NVFP4 at load " + f"(FC1 amax bound {fc1_amax:.6g}, FC2 amax {fc2_amax:.6g})" + ) + expert_weights = build_nvfp4_expert_weights_from_bf16( + weights, + experts_prefix, + expert_start, + expert_end, + config.num_experts, + config.intermediate_size, + fc1_amax, + fc2_amax, + accounting, + ) + # Fast paths tuned from served routing of the checkpoint's quantized layers + # (hot_replicas' census) have no data for this layer. + moe.experts_quantized_at_load = True + else: + expert_weights = build_expert_weights( + weights, + experts_prefix, + expert_start, + expert_end, + config.num_experts, + config.intermediate_size, + is_quantized, + accounting, + ) moe.experts.load_weights([expert_weights]) expert_weights.clear() diff --git a/tensorrt_llm/_torch/models/modeling_musespark_open.py b/tensorrt_llm/_torch/models/modeling_musespark_open.py index 0c2dd676ddec..6068e8e03475 100644 --- a/tensorrt_llm/_torch/models/modeling_musespark_open.py +++ b/tensorrt_llm/_torch/models/modeling_musespark_open.py @@ -84,7 +84,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping -from ...models.modeling_utils import QuantConfig +from ...models.modeling_utils import QuantAlgo, QuantConfig from ..attention.backends import AttentionMetadata from ..attention.backends.interface import ( PositionalEmbeddingParams, @@ -865,6 +865,62 @@ def _fused_pre_expert_quant_enabled() -> bool: ) +def excluded_layer_experts_nvfp4_enabled() -> bool: + """Whether the checkpoint's quant-excluded MoE layers run NVFP4 experts. + + Off unless MUSESPARK_OPEN_L061_EXPERTS_NVFP4 is ``1``. When on, the experts + of every layer the checkpoint excludes whole (``layers.0.`` / + ``layers.61.``) are built on the global NVFP4 leaf and the loader quantizes + their BF16 checkpoint tensors at load + (``musespark_open_weight_mapper.build_nvfp4_expert_weights_from_bf16``). + Attention in those layers stays BF16 through its own ``self_attn`` + exclusion. MUSESPARK_OPEN_L061_EXPERTS_NVFP4_LAYERS (comma-separated layer + indices) limits it to some of those layers. Read at model construction, + before any module is built. + """ + return os.environ.get("MUSESPARK_OPEN_L061_EXPERTS_NVFP4", "0").strip() == "1" + + +def unexclude_layer_experts( + quant_config: QuantConfig, num_layers: int, layers: Optional[List[int]] = None +) -> List[int]: + """Drop the exclusions that keep whole layers' experts out of quantization. + + Removes every ``exclude_modules`` entry that matches the + ``model.layers.N.feed_forward.experts`` of a layer in ``layers`` (every + layer when None), so those experts resolve to the global quantization + format. Returns the indices of the layers whose experts changed format. + + Raises: + ValueError: if a dropped entry was also what excluded some layer's + attention, which must stay in the format the checkpoint stores, or + the experts of an excluded layer outside ``layers``. + """ + patterns = list(quant_config.exclude_modules or []) + selected = set(range(num_layers) if layers is None else layers) + + def matches(entries: List[str], name: str) -> bool: + return QuantConfig(exclude_modules=entries).is_module_excluded_from_quantization(name) + + expert_names = [f"model.layers.{index}.feed_forward.experts" for index in range(num_layers)] + attention_names = [f"model.layers.{index}.self_attn" for index in range(num_layers)] + targets = [name for index, name in enumerate(expert_names) if index in selected] + before = [matches(patterns, name) for name in expert_names] + attention_before = [matches(patterns, name) for name in attention_names] + kept = [entry for entry in patterns if not any(matches([entry], n) for n in targets)] + after = [matches(kept, name) for name in expert_names] + changed = [index for index in range(num_layers) if before[index] and not after[index]] + attention_after = [matches(kept, name) for name in attention_names] + if attention_after != attention_before or any(index not in selected for index in changed): + raise ValueError( + "MuseSparkOpen: dropping the layer exclusions " + f"{sorted(set(patterns) - set(kept))} would change the format of more than " + f"the experts of layers {sorted(selected)}" + ) + quant_config.exclude_modules = kept + return changed + + @triton.jit def _routing_weights_kernel( scores_ptr, @@ -2942,6 +2998,22 @@ def __init__(self, model_config: ModelConfig[MuseSparkOpenConfig]): if quant_config is not None and quant_config.exclude_modules: quant_config.exclude_modules = normalize_exclude_modules(quant_config.exclude_modules) config = model_config.pretrained_config + if ( + quant_config is not None + and quant_config.quant_algo == QuantAlgo.NVFP4 + and quant_config.exclude_modules + and excluded_layer_experts_nvfp4_enabled() + ): + only = os.environ.get("MUSESPARK_OPEN_L061_EXPERTS_NVFP4_LAYERS", "").strip() + layers = unexclude_layer_experts( + quant_config, + config.num_hidden_layers, + [int(item) for item in only.split(",") if item.strip()] if only else None, + ) + logger.info( + f"MuseSparkOpen: experts of layers {layers} run NVFP4, quantized from the " + f"BF16 checkpoint at load; exclude_modules is now {quant_config.exclude_modules}" + ) assert model_config.mapping.pp_size == 1, ( "MuseSparkOpen does not support pipeline parallelism" ) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_l061_nvfp4.py b/tests/unittest/_torch/modeling/test_musespark_open_l061_nvfp4.py new file mode 100644 index 000000000000..f980a811f0cb --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_l061_nvfp4.py @@ -0,0 +1,395 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""NVFP4 experts for the layers the checkpoint stores in BF16 (MUSESPARK_OPEN_L061_EXPERTS_NVFP4). + +The CPU tests pin the exclusion rewrite and the FC1 activation bound. The GPU +tests check the load-time quantizer against the dequantization the checkpoint's +quantized layers use, and the whole NVFP4 MoE layer against the BF16 one it +replaces. With MSO_CHECKPOINT_DIR (the NVFP4 checkpoint) they use the real +layer 0 / 61 experts; with MSO_L061_DUMP_DIR as well (saved expert inputs: files +holding ``expert_input`` and ``router_logits``) they use real expert inputs and report +the FC2-input amax those inputs reach. +""" + +import glob +import json +import os + +import pytest +import torch + +from tensorrt_llm._torch.models.checkpoints.hf.musespark_open_weight_mapper import ( + NVFP4_GLOBAL_RANGE, + WeightAccounting, + build_expert_weights, + build_nvfp4_expert_weights_from_bf16, + nvfp4_quantize_weight, + pre_expert_norm_amax, +) +from tensorrt_llm._torch.models.modeling_musespark_open import ( + normalize_exclude_modules, + unexclude_layer_experts, +) +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + +HIDDEN, INTERMEDIATE, EXPERTS, TOP_K = 4096, 4096, 256, 8 +CHECKPOINT = os.environ.get("MSO_CHECKPOINT_DIR", "") +DUMPS = os.environ.get("MSO_L061_DUMP_DIR", "") +_PREFIX = "model.language_model.layers.{layer}.mlp." + + +def _checkpoint_reader(): + from safetensors import safe_open + + index = json.load(open(os.path.join(CHECKPOINT, "model.safetensors.index.json"))) + weight_map = index["weight_map"] + + def read(key): + with safe_open(os.path.join(CHECKPOINT, weight_map[key]), framework="pt") as handle: + return handle.get_tensor(key) + + return read + + +def test_unexclude_layer_experts_keeps_attention_excluded(): + config = QuantConfig( + quant_algo=QuantAlgo.NVFP4, + exclude_modules=normalize_exclude_modules( + ["layers.0.", "layers.61.", "self_attn", "lm_head"] + ), + ) + assert config.is_module_excluded_from_quantization("model.layers.0.feed_forward.experts") + assert unexclude_layer_experts(config, 62) == [0, 61] + assert config.exclude_modules == ["*self_attn*", "lm_head"] + for layer in (0, 1, 61): + assert not config.is_module_excluded_from_quantization( + f"model.layers.{layer}.feed_forward.experts" + ) + assert config.is_module_excluded_from_quantization(f"model.layers.{layer}.self_attn") + assert unexclude_layer_experts(config, 62) == [] + + +def test_unexclude_layer_experts_one_layer(): + config = QuantConfig( + quant_algo=QuantAlgo.NVFP4, + exclude_modules=["*layers.0.*", "*layers.61.*", "*self_attn*", "lm_head"], + ) + assert unexclude_layer_experts(config, 62, [61]) == [61] + assert config.exclude_modules == ["*layers.0.*", "*self_attn*", "lm_head"] + assert config.is_module_excluded_from_quantization("model.layers.0.feed_forward.experts") + + +def test_unexclude_layer_experts_refuses_to_requantize_attention(): + config = QuantConfig(quant_algo=QuantAlgo.NVFP4, exclude_modules=["*layers.0.*", "lm_head"]) + with pytest.raises(ValueError, match="format of more than"): + unexclude_layer_experts(config, 62) + + +def test_pre_expert_norm_amax_bounds_every_input(): + torch.manual_seed(5) + weight = torch.randn(HIDDEN).to(torch.bfloat16) + bound = pre_expert_norm_amax(weight, HIDDEN) + gain = (weight.float() + 1.0).to(torch.bfloat16).float() + x = torch.randn(64, HIDDEN) * torch.rand(64, 1) * 10 + x[0] = 0.0 + x[0, int(gain.abs().argmax())] = 300.0 + y = x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + 1e-5) * gain + assert y.abs().max() <= bound + # A one-hot input on the largest gain reaches the bound. + assert y[0].abs().max() >= bound * 0.999 + + +def _dequantize(packed, block_scale, weight_scale_2): + return torch.ops.tensorrt_llm.e2m1_and_ufp8sf_scale_to_float_v2( + packed.cpu(), + block_scale.view(torch.uint8).flatten().cpu(), + weight_scale_2.float().cpu().reshape(1), + 16, + 1, + False, + ) + + +_E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] +_E2M1 = torch.tensor(_E2M1_MAGNITUDES + [-v for v in _E2M1_MAGNITUDES]) + + +def _dequantize_cuda(packed, block_scale, weight_scale_2): + """GPU dequantization of :func:`nvfp4_quantize_weight`'s output (low nibble first).""" + packed = packed.cuda() + table = _E2M1.cuda() + values = torch.stack((table[(packed & 0xF).long()], table[(packed >> 4).long()]), dim=-1) + values = values.reshape(packed.shape[0], -1) + scale = block_scale.cuda().float().repeat_interleave(16, dim=-1) * weight_scale_2.float().cuda() + return values * scale + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_nvfp4_quantize_weight_round_trip(): + torch.manual_seed(7) + weight = (torch.randn(512, HIDDEN) * 0.02).to(torch.bfloat16) + weight[3, 17] = 0.9 + packed, block_scale, weight_scale_2 = nvfp4_quantize_weight(weight) + assert packed.dtype == torch.uint8 and packed.shape == (512, HIDDEN // 2) + assert block_scale.dtype == torch.float8_e4m3fn and block_scale.shape == (512, HIDDEN // 16) + assert weight_scale_2.shape == () and weight_scale_2.dtype == torch.float32 + torch.testing.assert_close( + weight_scale_2, weight[3, 17].float() / NVFP4_GLOBAL_RANGE, rtol=1e-6, atol=0.0 + ) + restored = _dequantize(packed, block_scale, weight_scale_2) + relative = (restored - weight.float()).norm() / weight.float().norm() + print(f"random weight round trip relative_l2={relative.item():.4f}") + assert relative < 0.15 + # The amax element maps onto E2M1 6 x E4M3 448 exactly. + assert abs(restored[3, 17].item() / weight[3, 17].float().item() - 1.0) < 1e-6 + torch.testing.assert_close( + _dequantize_cuda(packed, block_scale, weight_scale_2).cpu(), restored, rtol=0.0, atol=0.0 + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not CHECKPOINT, reason="CUDA and MSO_CHECKPOINT_DIR required" +) +def test_nvfp4_quantize_weight_matches_checkpoint_layout(): + """Re-quantizing a dequantized checkpoint weight reproduces it. + + The values of a dequantized NVFP4 tensor are on that tensor's own grid, so a + quantizer with the same nibble order, scale layout and rounding maps them + back almost exactly; a swapped nibble order or a transposed scale layout + would not. + """ + read = _checkpoint_reader() + prefix = _PREFIX.format(layer=1) + "experts." + packed = read(prefix + "gate_up_proj")[5, :INTERMEDIATE] + block_scale = read(prefix + "gate_up_proj_weight_scale")[5, :INTERMEDIATE] + weight_scale_2 = read(prefix + "gate_up_proj_weight_scale_2")[5, 0] + original = _dequantize(packed, block_scale, weight_scale_2) + again = nvfp4_quantize_weight(original.to(torch.bfloat16)) + restored = _dequantize(*again) + relative = (restored - original).norm() / original.norm() + same_bytes = (again[0].cpu() == packed).float().mean() + print( + f"layer 1 expert 5 w1 re-quantized: relative_l2={relative.item():.3g} " + f"identical_bytes={same_bytes.item():.4f}" + ) + assert relative < 2e-2 + assert same_bytes > 0.9 + + +def _layer_tensors(read, layer): + mlp = _PREFIX.format(layer=layer) + return { + "w13": read(mlp + "experts.gate_up_proj"), + "w2": read(mlp + "experts.down_proj"), + "post_norm": read(mlp + "experts.post_expert_norm.weight"), + "pre_norm": read(mlp + "pre_expert_norm.weight"), + "bias": read(mlp + "gate.e_score_correction_bias"), + } + + +def _inputs(layer, gain, count): + """Real (expert input, router logits) calls of ``layer``, else synthetic ones.""" + calls = [] + for path in sorted(glob.glob(os.path.join(DUMPS, f"layer{layer}_call*.pt")))[:count]: + record = torch.load(path) + calls.append((record["expert_input"], record["router_logits"], os.path.basename(path))) + if calls or DUMPS: + return calls + torch.manual_seed(11 + layer) + x = torch.randn(2048, HIDDEN) + x = x / torch.sqrt(x.pow(2).mean(-1, keepdim=True)) * gain + return [(x.to(torch.bfloat16), torch.randn(2048, EXPERTS) * 2, "synthetic")] + + +def _make_moe(bias, quantized): + from transformers import PretrainedConfig + + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_musespark_open import MuseSparkOpenRoutingMethod + from tensorrt_llm._torch.moe.fused_moe import create_moe + from tensorrt_llm.mapping import Mapping + + config = ModelConfig( + pretrained_config=PretrainedConfig( + hidden_size=HIDDEN, + intermediate_size=INTERMEDIATE, + num_experts=EXPERTS, + torch_dtype=torch.bfloat16, + ), + mapping=Mapping(), + moe_backend="TRTLLM", + ) + routing = MuseSparkOpenRoutingMethod( + top_k=TOP_K, + num_experts=EXPERTS, + callable_e_score_correction_bias=lambda: bias, + renorm_eps=1e-15, + ) + return create_moe( + routing_method=routing, + num_experts=EXPERTS, + hidden_size=HIDDEN, + intermediate_size=INTERMEDIATE, + dtype=torch.bfloat16, + reduce_results=False, + model_config=config, + override_quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4) if quantized else None, + layer_idx=0, + ).cuda() + + +def _fc2_input_amax(w13, x, selected): + """max |silu(x w1^T) * (x w3^T)| over the (token, selected expert) pairs.""" + amax = torch.zeros(EXPERTS) + for expert in torch.unique(selected).tolist(): + rows = (selected == expert).any(-1).nonzero().flatten() + h = x[rows] @ w13[expert].T + act = torch.nn.functional.silu(h[:, :INTERMEDIATE].float()) * h[:, INTERMEDIATE:].float() + amax[expert] = act.abs().max().item() + return amax + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not CHECKPOINT, reason="CUDA and MSO_CHECKPOINT_DIR required" +) +@pytest.mark.parametrize("layer", [0, 61]) +@torch.inference_mode() +def test_excluded_layer_nvfp4_moe_matches_bf16(layer): + """The NVFP4 MoE built from the BF16 experts, against the BF16 MoE, on one GPU.""" + from tensorrt_llm._torch.models.modeling_musespark_open import post_expert_norm_combine + + read = _checkpoint_reader() + tensors = _layer_tensors(read, layer) + gain = (tensors["pre_norm"].float() + 1.0).to(torch.bfloat16).float() + fc1_amax = pre_expert_norm_amax(tensors["pre_norm"], HIDDEN) + calls = _inputs(layer, gain, int(os.environ.get("MSO_L061_MAX_CALLS", "64"))) + assert calls, f"no layer {layer} dumps under {DUMPS}" + bias = tensors["bias"].float().cuda() + post_norm = tensors["post_norm"].cuda() + checkpoint = {"e.w13_weight": tensors["w13"], "e.w2_weight": tensors["w2"]} + + bf16 = _make_moe(bias, quantized=False) + bf16.load_weights( + [ + build_expert_weights( + checkpoint, "e", 0, EXPERTS, EXPERTS, INTERMEDIATE, False, WeightAccounting() + ) + ] + ) + bf16.post_load_weights() + + # Calibration: the FC2-input amax these inputs reach, before choosing it. + w13_gpu = tensors["w13"].cuda() + observed_fc1, observed_fc2 = 0.0, torch.zeros(EXPERTS) + for x, logits, name in calls: + x, logits = x.cuda(), logits.cuda() + selected, _ = bf16.routing_method.apply(logits) + per_expert = _fc2_input_amax(w13_gpu, x, selected) + observed_fc2 = torch.maximum(observed_fc2, per_expert) + observed_fc1 = max(observed_fc1, x.float().abs().max().item()) + distinct = torch.unique(x[:, :64].float(), dim=0).shape[0] / x.shape[0] + print( + f"layer {layer} {name}: rows={x.shape[0]} distinct_rows={distinct:.3f} " + f"fc1_amax={x.float().abs().max().item():.4g} fc2_amax={per_expert.max().item():.4g}" + ) + print( + f"layer {layer} calibration over {len(calls)} calls: fc1 observed {observed_fc1:.4g} " + f"(bound {fc1_amax:.4g}), fc2 observed max {observed_fc2.max().item():.4g}, " + f"p50/p90/p99 over experts {observed_fc2.quantile(0.5).item():.4g}/" + f"{observed_fc2.quantile(0.9).item():.4g}/{observed_fc2.quantile(0.99).item():.4g}" + ) + assert observed_fc1 <= fc1_amax + + configured = os.environ.get("MUSESPARK_OPEN_L061_EXPERTS_NVFP4_FC2_AMAX", "") + fc2_amax = {int(k): float(v) for k, v in (i.split("=") for i in configured.split(",") if i)} + fc2_amax = fc2_amax.get(layer, 2.0 * observed_fc2.max().item()) + print(f"layer {layer} FC2 amax used: {fc2_amax:.4g}") + assert fc2_amax >= observed_fc2.max().item() + del w13_gpu + + def nvfp4_moe(fc1): + moe = _make_moe(bias, quantized=True) + moe.load_weights( + [ + build_nvfp4_expert_weights_from_bf16( + checkpoint, + "e", + 0, + EXPERTS, + EXPERTS, + INTERMEDIATE, + fc1, + fc2_amax, + WeightAccounting(), + ) + ] + ) + moe.post_load_weights() + return moe + + # The served configuration, then two diagnostics: FC1 scaled from the observed amax + # instead of the bound, and NVFP4 weights with BF16 activations (weight error only). + variants = {"nvfp4": nvfp4_moe(fc1_amax)} + if os.environ.get("MSO_L061_DIAGNOSTICS", "1") == "1": + variants["nvfp4_fc1_observed"] = nvfp4_moe(2.0 * observed_fc1) + dequantized = {} + for expert in range(EXPERTS): + for name, weight in ( + ("w1", tensors["w13"][expert, :INTERMEDIATE]), + ("w3", tensors["w13"][expert, INTERMEDIATE:]), + ("w2", tensors["w2"][expert]), + ): + dequantized[f"{expert}.{name}.weight"] = _dequantize_cuda( + *nvfp4_quantize_weight(weight) + ).to(torch.bfloat16) + weight_only = _make_moe(bias, quantized=False) + weight_only.load_weights([dequantized]) + weight_only.post_load_weights() + del dequantized + variants["weights_only"] = weight_only + + worst = 1.0 + for x, logits, name in calls[: int(os.environ.get("MSO_L061_COMPARE_CALLS", "8"))]: + x, logits = x.cuda(), logits.cuda() + tokens = x.shape[0] + _, routing = bf16.routing_method.apply(logits) + + def run(moe): + gemm2, _, index = moe(x, logits, do_finalize=False) + return post_expert_norm_combine( + gemm2, + index, + routing, + post_norm, + 1e-8, + num_tokens=tokens, + top_k=TOP_K, + hidden_size=HIDDEN, + ).float() + + reference = run(bf16) + for label, moe in variants.items(): + actual = run(moe) + relative = (actual - reference).norm() / reference.norm() + cosine = torch.nn.functional.cosine_similarity(actual, reference, dim=-1) + print( + f"layer {layer} {name}: {label} vs BF16 relative_l2={relative.item():.4f} " + f"token cosine mean={cosine.mean().item():.5f} " + f"p1={cosine.quantile(0.01).item():.5f} min={cosine.min().item():.5f}" + ) + assert torch.isfinite(actual).all() + if label == "nvfp4": + worst = min(worst, cosine.mean().item()) + assert worst > float(os.environ.get("MSO_L061_MIN_COSINE", "0.97")) From bc3e8c372d4f5a580169fea796f0a834c720c6b9 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:10:58 +0000 Subject: [PATCH 06/15] [None][feat] modeling_v2 musespark_open tep: FP8 copies of the dense projections for large context calls New kernels/context_fp8_dense: at load, FP8 copies of the qkv, o_proj, pool and unpool projection weights; at run time, context calls of at least min_rows rows (1024 and above) quantize their input with a one-pass MXFP8 quant and run the FP8 GEMM with fixed Rubin tactics, plus an optional MXFP8 attention-input gather on sequence-parallel steps. Decode steps never reach the row floor. The install refuses unless the device is sm107, the Rubin CuTe DSL is available and the projections are plain bias-free BF16. This is a precision change. The lever resolver keeps it off unless it is requested explicitly. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../kernels/context_fp8_dense/__init__.py | 3 + .../kernels/context_fp8_dense/install.py | 394 ++++++++++++++++++ .../kernels/context_fp8_dense/quant.py | 116 ++++++ .../test_musespark_open_context_fp8_dense.py | 166 ++++++++ 4 files changed, 679 insertions(+) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/__init__.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/install.py create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/quant.py create mode 100644 tests/unittest/_torch/modeling/test_musespark_open_context_fp8_dense.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/__init__.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/__init__.py new file mode 100644 index 000000000000..68db3850477b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""context_fp8_dense: the context steps' dense projections on the Rubin MXFP8 GEMM (a lossy lever).""" diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/install.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/install.py new file mode 100644 index 000000000000..a6146272428b --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/install.py @@ -0,0 +1,394 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``context_fp8_dense``: the dense projections of large context calls on the Rubin MXFP8 GEMM (sm107). + +A lossy lever (``levers.RULES``: on only when its knob sets it): MXFP8, E4M3 values with +one UE8M0 scale per 1x32 block of weights and of activations. The plan gives ``install`` +its ``MIN_ROWS`` and ``PROJECTIONS`` parameters. + +Every decoder layer's QKV projection, output projection, pooler and unpooler (the BF16 +``Linear`` modules ``PROJECTIONS`` names) keeps its BF16 weight, which every decode path +and every small call still reads, and gets an MXFP8 copy of it. A call through the +module's quant method (``Linear.forward``) with at least ``min_rows`` rows -- the context +steps: the sequence-parallel arm's gathered attention rows and its shard of the MoE rows, +and the replicated forward of a large context step -- quantizes its BF16 activation +(``quant.mxfp8_quantize``, one pass at HBM speed; the same bits as +``trtllm::mxfp8_quantize``) and runs the Rubin CuTe DSL MXFP8 GEMM +(``CuteDSLMXFP8RubinLinear``, the runner behind ``trtllm::cute_dsl_mxfp8_gemm_rubin``) at +the tactic ``TACTICS`` lists for its weight shape and rows; every other call runs the BF16 +GEMM unchanged. The tactics are fixed, so no server start compiles or profiles the +runner's few hundred tactics; ``install`` compiles the listed ones before any capture. +Decode steps (at most ``max_batch_size * (draft_len + 1)`` rows, far below ``min_rows``) +and every kernel that reads ``module.weight`` directly are unchanged. + +The quant method stays an ``UnquantizedLinearMethod`` (a subclass), so every check that +gates a fast path on a BF16 projection still passes and still reads the BF16 weight. The +copy-engine gather of the attention input (``context_sequence_parallel.ce_gather``) +multiplies the QKV row blocks with that BF16 weight itself; where it is wired +(``SHARD_PARTS_BY_SM``) the QKV projection keeps BF16 on every call, so a step never mixes +the two, and the install says so. + +With ``context_fp8_attn_gather`` (``gather``) the sequence-parallel arm's attention +input is quantized on the shard, before its all-gather (``Fp8AttentionInputGather``, the +decoder layer's ``_sp_attn_input_gather`` hook): the gather moves E4M3 bytes and row-major +scales instead of BF16, and the QKV projection consumes the gathered ``GatheredMxfp8`` +rows. Every block of 32 lies in one row, so the values and scales are the bytes the +quantization after the gather produces. + +Numerics: lossy (MXFP8 activations and weights on the context calls); gated by accuracy +runs, never on by default. The copies cost 1.03 bytes per weight element on top of the +BF16 weights (``install`` logs the total), which the KV cache sizing after the load sees. + +Install after the weights are final and before the KV cache is sized and any CUDA graph +is captured. Agreed over the TP group (``agreed_install``). +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import torch + +from tensorrt_llm._torch.modules.linear import Linear, UnquantizedLinearMethod +from tensorrt_llm.logger import logger + +from ... import dispatch as _tp_agreement +from . import quant as _quant + +#: The log name of this path. +NAME = "context_fp8_dense" +#: The projections a copy can be installed on (the lever's ``PROJECTIONS`` choices). +PROJECTIONS = ("qkv", "o_proj", "pool", "unpool") +#: The SM version this path is written for: the GEMM is the Rubin (sm107) CuTe DSL one. +SERVED_SM = 107 +#: The fewest rows of a call that runs FP8 by default (the lever's ``MIN_ROWS`` default). +DEFAULT_MIN_ROWS = 1024 + +_BASE_256_21_CLC = ( + "base", + (256, 256, 128), + (256, 256, 64), + (2, 1), + False, + False, + "clc_dynamic", + "m", + 1, +) +_BASE_256_22_CLC = ( + "base", + (256, 256, 128), + (256, 256, 64), + (2, 2), + False, + False, + "clc_dynamic", + "m", + 1, +) +_BASE_256_42 = ("base", (256, 256, 128), (256, 256, 64), (4, 2), False, False, "static", "m", 1) +_BASE_512_21 = ("base", (512, 256, 128), (256, 256, 64), (2, 1), False, False, "static", "m", 1) +_BASE_512_22_N = ("base", (512, 256, 128), (256, 256, 64), (2, 2), False, False, "static", "n", 1) +_MIXED_256 = ( + "mixed_clusters", + (256, 256, 128), + (256, 256, 64), + (4, 2), + (2, 1), + False, + False, + "static", + "m", +) +_MIXED_512 = ( + "mixed_clusters", + (512, 256, 128), + (256, 256, 64), + (4, 2), + (2, 1), + False, + False, + "static", + "m", +) + +#: ``(out_features, in_features) -> ((fewest rows, tactic), ...)``, rows ascending: a call of M +#: rows runs the tactic of the last entry with ``fewest rows <= M``. The fastest non-swapped +#: tactic per row count of a sweep on one Rubin GPU under CUDA-graph replay (this model's TP4 shapes). +TACTICS: Dict[Tuple[int, int], Tuple[Tuple[int, Tuple[Any, ...]], ...]] = { + (4608, 8192): ( # QKV + (0, _BASE_256_22_CLC), + (1536, _BASE_512_21), + (2048, _BASE_512_22_N), + (3072, _MIXED_256), + (6144, _BASE_256_22_CLC), + (8192, _MIXED_512), + ), + (8192, 2048): ( # o_proj + (0, _BASE_512_21), + (1536, _BASE_256_42), + (3072, _BASE_256_21_CLC), + ), + (4096, 8192): ( # pool + (0, _MIXED_256), + (2048, _BASE_512_22_N), + (3072, _MIXED_512), + (4096, _BASE_256_42), + (6144, _BASE_512_21), + (8192, _BASE_256_21_CLC), + ), + (8192, 4096): ( # unpool + (0, _BASE_512_21), + (1536, _BASE_256_42), + (3072, _MIXED_512), + (4096, _MIXED_256), + (8192, _BASE_512_21), + ), +} +#: Any other weight shape: the runner's default tactic. +DEFAULT_TACTIC = -1 + + +def _modules_of(layer) -> Dict[str, Linear]: + attention, moe = layer.self_attn, layer.feed_forward + return { + "qkv": attention.qkv_proj, + "o_proj": attention.o_proj, + "pool": getattr(moe, "pooler", None), + "unpool": getattr(moe, "unpooler", None), + } + + +def _runner(): + from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops + + return cute_dsl_custom_ops.CuteDSLMXFP8RubinLinear( + output_dtype=torch.bfloat16, use_tvm_ffi=True + ), cute_dsl_custom_ops._get_mxfp8_gemm_alpha + + +class Fp8Projection: + """One projection's MXFP8 weight copy and its GEMM: ``(rows, K)`` BF16 in, ``(rows, N)`` BF16 out.""" + + def __init__(self, weight: torch.Tensor, min_rows: int) -> None: + if weight.dim() != 2 or weight.dtype != torch.bfloat16 or weight.shape[1] % _quant.BLOCK_K: + raise ValueError( + f"needs a 2-D BF16 weight with K a multiple of {_quant.BLOCK_K}, got " + f"{weight.dtype} {tuple(weight.shape)}" + ) + self.min_rows = min_rows + self.out_features, self.in_features = weight.shape + self.tactics: Sequence[Tuple[int, Any]] = TACTICS.get( + (self.out_features, self.in_features), ((0, DEFAULT_TACTIC),) + ) + self._runner, get_alpha = _runner() + self._alpha = get_alpha(weight.device) + with torch.no_grad(): + self.weight, self.scale = _quant.mxfp8_quantize(weight.contiguous()) + + def nbytes(self) -> int: + return self.weight.numel() * self.weight.element_size() + self.scale.numel() + + def tactic(self, rows: int) -> Any: + chosen = self.tactics[0][1] + for fewest, tactic in self.tactics: + if rows >= fewest: + chosen = tactic + return chosen + + def takes(self, input: torch.Tensor, bias: Optional[torch.Tensor]) -> bool: + return ( + bias is None + and input.dim() == 2 + and input.shape[0] >= self.min_rows + and input.shape[1] == self.in_features + and input.dtype == torch.bfloat16 + and input.stride(1) == 1 + and not torch.compiler.is_compiling() + ) + + def __call__(self, input: torch.Tensor, tactic: Any = None) -> torch.Tensor: + values, scale = _quant.mxfp8_quantize(input) + return self.project(values, scale, tactic) + + def project( + self, values: torch.Tensor, scale: torch.Tensor, tactic: Any = None + ) -> torch.Tensor: + """The GEMM of already quantized rows (E4M3 ``[M, K]``, 128x4-swizzled scales).""" + if tactic is None: + tactic = self.tactic(values.shape[0]) + return self._runner([values, self.weight, scale, self.scale, self._alpha], tactic=tactic) + + +class GatheredMxfp8: + """A step's attention input rows, gathered in MXFP8 (``Fp8AttentionInputGather``). + + Its only consumer is the QKV projection whose quant method is ``ContextFp8LinearMethod``. + """ + + def __init__(self, values: torch.Tensor, scale: torch.Tensor) -> None: + self.values = values + self.scale = scale + + @property + def shape(self) -> torch.Size: + return self.values.shape + + +class Fp8AttentionInputGather: + """The decoder layer's ``_sp_attn_input_gather``: quantize the shard, then all-gather it. + + Steps below ``min_rows`` rows run the NCCL all-gather of the BF16 rows the layer runs + without the hook. + """ + + def __init__(self, mapping, min_rows: int) -> None: + self.mapping = mapping + self.min_rows = min_rows + + def __call__(self, attn_input: torch.Tensor, rows: int): + from tensorrt_llm._torch.distributed import allgather + + if ( + rows < self.min_rows + or attn_input.dim() != 2 + or attn_input.dtype != torch.bfloat16 + or attn_input.stride(1) != 1 + or attn_input.shape[1] % _quant.BLOCK_K + or torch.compiler.is_compiling() + ): + return allgather(attn_input, self.mapping, dim=0)[:rows] + values, scales = _quant.mxfp8_quantize(attn_input, swizzled=False) + values, scales = allgather([values.view(torch.uint8), scales], self.mapping, dim=0) + return GatheredMxfp8( + values[:rows].view(torch.float8_e4m3fn), _quant.swizzle_scales(scales[:rows]) + ) + + +class ContextFp8LinearMethod(UnquantizedLinearMethod): + """``UnquantizedLinearMethod`` whose calls of at least the copy's ``min_rows`` rows run MXFP8.""" + + def apply(self, module: Linear, input: torch.Tensor, bias: Optional[torch.Tensor]): + fp8 = module.__dict__.get("_context_fp8") + if isinstance(input, GatheredMxfp8): + if fp8 is None or bias is not None: + raise RuntimeError( + "context_fp8_dense: gathered MXFP8 rows reached a projection without an FP8 copy" + ) + return fp8.project(input.values, input.scale) + if fp8 is not None and fp8.takes(input, bias): + return fp8(input) + return super().apply(module, input, bias) + + +def _ce_gather_wired(sm: Optional[int]) -> bool: + from ..context_sequence_parallel.forward import SHARD_PARTS_BY_SM + + return "ce_gather" in SHARD_PARTS_BY_SM.get(sm, ()) + + +def install( + text_model, + mapping, + compute_capability: Optional[Tuple[int, int]], + min_rows: int = DEFAULT_MIN_ROWS, + projections: Sequence[str] = PROJECTIONS, + gather: bool = False, +) -> List[str]: + """Give ``projections`` of every layer an MXFP8 copy; returns the projections installed. + + Calls of ``min_rows`` rows or more run FP8; with ``gather`` (``context_fp8_attn_gather``) + the sequence-parallel attention input is quantized before its all-gather. Collective over + the TP group (the outcome is agreed): raises ValueError on every rank alike when some rank + cannot serve the request (not sm107, no Rubin CuTe DSL GEMM, a projection that is not a + bias-free BF16 ``UnquantizedLinearMethod`` Linear, or a failing first call). A call with + LoRA parameters runs the copy as LoRA's base GEMM, as it runs the BF16 one. + """ + sm = _tp_agreement.sm_of(compute_capability) + + def build() -> List[str]: + if sm != SERVED_SM: + raise ValueError(f"[musespark_open] {NAME}: written for sm{SERVED_SM}, got sm{sm}") + from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + + if not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + raise ValueError( + f"[musespark_open] {NAME}: the Rubin CuTe DSL package is not available" + ) + unknown = sorted(set(projections) - set(PROJECTIONS)) + if unknown or min_rows < 1: + raise ValueError( + f"[musespark_open] {NAME}: projections {list(projections)} of {list(PROJECTIONS)} " + f"from {min_rows} rows" + ) + names = [name for name in PROJECTIONS if name in projections] + if "qkv" in names and _ce_gather_wired(sm): + names.remove("qkv") + logger.warning( + f"[musespark_open] {NAME}: the copy-engine gather is wired on sm{sm} and multiplies " + "the QKV row blocks with the BF16 weight, so the QKV projection keeps BF16" + ) + sites = [] + for layer in text_model.layers: + modules = _modules_of(layer) + for name in names: + module = modules[name] + if ( + not isinstance(module, Linear) + or type(module.quant_method) is not UnquantizedLinearMethod + ): + raise ValueError( + f"[musespark_open] {NAME}: layer {layer.layer_idx} {name} is not a built-in " + f"unquantized Linear ({type(getattr(module, 'quant_method', None)).__name__})" + ) + if module.bias is not None: + raise ValueError( + f"[musespark_open] {NAME}: layer {layer.layer_idx} {name} has a bias" + ) + sites.append(module) + added = 0 + try: + copies = [] + for module in sites: + copy = Fp8Projection(module.weight.data, min_rows) + copies.append((module, copy)) + added += copy.nbytes() + # Compile every listed tactic of every weight shape (and create the GEMM's alpha) + # before any CUDA graph capture. + seen = set() + for module, copy in copies: + key = (copy.out_features, copy.in_features) + if key in seen: + continue + seen.add(key) + for fewest, tactic in copy.tactics: + rows = max(fewest, min_rows) + probe = torch.zeros( + (rows, copy.in_features), dtype=torch.bfloat16, device=module.weight.device + ) + copy(probe, tactic) + torch.cuda.synchronize() + torch.cuda.empty_cache() + except Exception as error: # noqa: BLE001 - the failure must reach the agreement + raise RuntimeError( + f"[musespark_open] {NAME}: could not build the MXFP8 copies: {error!r}" + ) from error + for module, copy in copies: + fp8_method = ContextFp8LinearMethod.__new__(ContextFp8LinearMethod) + fp8_method.__dict__.update(module.quant_method.__dict__) + module.quant_method = fp8_method + module._context_fp8 = copy + gathers = 0 + if "qkv" in names and gather: + for layer in text_model.layers: + if getattr(layer, "_sp_attn_input_gather", None) is None: + layer._sp_attn_input_gather = Fp8AttentionInputGather(layer.mapping, min_rows) + gathers += 1 + logger.info( + f"[musespark_open] {NAME}: {names} of {len(text_model.layers)} layers run their calls of " + f">= {min_rows} rows on the Rubin MXFP8 GEMM; the copies add {added / 2**30:.2f} GiB per rank; " + f"{gathers} layers gather their attention input in MXFP8" + ) + return names + + return _tp_agreement.agreed_install(mapping, build, (ValueError, RuntimeError)) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/quant.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/quant.py new file mode 100644 index 000000000000..9f9e9a8366db --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_fp8_dense/quant.py @@ -0,0 +1,116 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 activation quantization for the Rubin MXFP8 GEMM, one Triton pass over the BF16 rows. + +``mxfp8_quantize(x)`` returns what ``trtllm::mxfp8_quantize(x, True)`` returns for a 2-D BF16 +``x`` whose K is a multiple of 128: E4M3 values ``[M, K]`` and one UE8M0 scale per 1x32 +block in the 128x4-swizzled layout (``pad_up(M, 128) * K / 32`` bytes, +``get_sf_out_offset_128x4`` in ``cpp/tensorrt_llm/kernels/quantization.cuh``), the scale +rows past M zero. Same rounding: scale = amax / 448 rounded up to a power of two, values +``x / scale`` rounded to nearest E4M3. Each program reads a 32 x 256 tile once, so the pass +runs at the bandwidth of one BF16 read and one E4M3 write. +""" + +from __future__ import annotations + +from typing import Tuple + +import torch +import triton +import triton.language as tl + +#: Rows and columns of one program's tile; K must be a multiple of BLOCK_K. +BLOCK_M = 32 +BLOCK_K = 256 + + +@triton.jit +def _mxfp8_quantize_kernel( + x_ptr, + q_ptr, + sf_ptr, + M, + stride_xm, + K: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, + SWIZZLED: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_k = tl.program_id(1) + rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + cols = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + in_rows = rows < M + x = tl.load( + x_ptr + rows[:, None].to(tl.int64) * stride_xm + cols[None, :], + mask=in_rows[:, None], + other=0.0, + ).to(tl.float32) + GROUPS: tl.constexpr = BLOCK_K // 32 + xg = tl.reshape(x, (BLOCK_M, GROUPS, 32)) + amax = tl.max(tl.abs(xg), axis=2) + bits = (amax * (1.0 / 448.0)).to(tl.int32, bitcast=True) + # UE8M0 rounded towards +inf: the exponent, plus one when any mantissa bit is set. + exponent = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32) + exponent = tl.minimum(exponent, 254) + # 2^(127 - exponent), the reciprocal of the scale; 0 for an all-zero block. + inverse = ((254 - exponent) << 23).to(tl.float32, bitcast=True) + inverse = tl.where(amax > 0.0, inverse, 0.0) + q = tl.reshape((xg * inverse[:, :, None]).to(tl.float8e4nv), (BLOCK_M, BLOCK_K)) + tl.store(q_ptr + rows[:, None].to(tl.int64) * K + cols[None, :], q, mask=in_rows[:, None]) + blocks = pid_k * GROUPS + tl.arange(0, GROUPS) + if SWIZZLED: + K_TILES: tl.constexpr = K // 128 + offsets = ( + (rows // 128).to(tl.int64)[:, None] * (K_TILES * 512) + + (blocks // 4)[None, :] * 512 + + (rows % 32)[:, None] * 16 + + ((rows % 128) // 32)[:, None] * 4 + + (blocks % 4)[None, :] + ) + tl.store(sf_ptr + offsets, exponent.to(tl.uint8)) + else: + offsets = rows.to(tl.int64)[:, None] * (K // 32) + blocks[None, :] + tl.store(sf_ptr + offsets, exponent.to(tl.uint8), mask=in_rows[:, None]) + + +def mxfp8_quantize(x: torch.Tensor, swizzled: bool = True) -> Tuple[torch.Tensor, torch.Tensor]: + """``(E4M3 [M, K], UE8M0 scales)`` of a 2-D BF16 ``x`` (last dim contiguous). + + The scales are 128x4-swizzled (``pad_up(M, 128) * K / 32`` bytes, what the GEMM reads), or + with ``swizzled=False`` row-major ``[M, K / 32]``, which row blocks of several calls can be + concatenated in (``swizzle_scales`` then lays them out for the GEMM). + """ + if x.dim() != 2 or x.dtype != torch.bfloat16 or x.stride(1) != 1: + raise ValueError( + f"needs a 2-D BF16 tensor with contiguous rows, got {x.dtype} {tuple(x.shape)}" + ) + m, k = x.shape + if k % BLOCK_K: + raise ValueError(f"K={k} must be a multiple of {BLOCK_K}") + padded_m = -(-m // 128) * 128 + q = torch.empty((m, k), dtype=torch.float8_e4m3fn, device=x.device) + if swizzled: + sf = torch.empty((padded_m * (k // 32),), dtype=torch.uint8, device=x.device) + grid = (padded_m // BLOCK_M, k // BLOCK_K) + else: + sf = torch.empty((m, k // 32), dtype=torch.uint8, device=x.device) + grid = (triton.cdiv(m, BLOCK_M), k // BLOCK_K) + _mxfp8_quantize_kernel[grid]( + x, + q, + sf, + m, + x.stride(0), + K=k, + BLOCK_M=BLOCK_M, + BLOCK_K=BLOCK_K, + SWIZZLED=swizzled, + num_warps=4, + ) + return q, sf + + +def swizzle_scales(scales: torch.Tensor) -> torch.Tensor: + """Row-major ``[M, K / 32]`` UE8M0 scales in the GEMM's 128x4-swizzled layout (rows past M zero).""" + return torch.ops.trtllm.block_scale_interleave(scales.contiguous()) diff --git a/tests/unittest/_torch/modeling/test_musespark_open_context_fp8_dense.py b/tests/unittest/_torch/modeling/test_musespark_open_context_fp8_dense.py new file mode 100644 index 000000000000..30722aa828ca --- /dev/null +++ b/tests/unittest/_torch/modeling/test_musespark_open_context_fp8_dense.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``context_fp8_dense``: the MXFP8 quant against the built-in op, the copy's GEMM against BF16, +and which calls take it (sm107).""" + +import importlib + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from tensorrt_llm._torch.modules.linear import Linear, UnquantizedLinearMethod + +_PACKAGE = ( + "tensorrt_llm._torch._experimental.modeling_v2.models.musespark_open." + "muse_spark_open_816b_a42b_nvfp4__sm_103__tep.kernels.context_fp8_dense" +) + +_ON_RUBIN = ( + torch.cuda.is_available() + and torch.cuda.get_device_capability() == (10, 7) + and IS_CUTLASS_DSL_RUBIN_AVAILABLE +) +needs_rubin = pytest.mark.skipif(not _ON_RUBIN, reason="needs sm107 with the Rubin CuTe DSL") + + +def _install(): + return importlib.import_module(f"{_PACKAGE}.install") + + +def _cos_diff(x: torch.Tensor, y: torch.Tensor) -> float: + x, y = x.double(), y.double() + return float(1 - 2 * (x * y).sum() / (x * x + y * y).sum()) + + +def test_lever_parameters_are_the_installers(): + """The lever table's ``context_fp8_dense`` parameters are the ones ``install`` takes.""" + install = _install() + levers = importlib.import_module(_PACKAGE.rsplit(".", 2)[0] + ".levers") + rule = levers.RULES["context_fp8_dense"] + assert rule.param("PROJECTIONS").choices == install.PROJECTIONS + assert rule.param("PROJECTIONS").default == install.PROJECTIONS + assert rule.param("MIN_ROWS").default == install.DEFAULT_MIN_ROWS + + +def test_tactic_table_is_ascending_and_unswapped(): + install = _install() + for shape, entries in install.TACTICS.items(): + rows = [fewest for fewest, _ in entries] + assert rows == sorted(rows) and rows[0] == 0, shape + for _, tactic in entries: + swap = tactic[4] if tactic[0] == "base" else tactic[5] + assert swap is False, (shape, tactic) + + +@needs_rubin +@pytest.mark.parametrize("m", [1, 100, 128, 1536, 4100, 8192]) +@pytest.mark.parametrize("k", [2048, 4096, 8192]) +def test_quant_matches_builtin_op(m, k): + quant = importlib.import_module(f"{_PACKAGE}.quant") + torch.manual_seed(m + k) + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + x *= 1 + 7 * torch.rand(m, 1, device="cuda", dtype=torch.bfloat16) + x[0, :64] = 0 # an all-zero block + values, scale = quant.mxfp8_quantize(x) + ref_values, ref_scale = torch.ops.trtllm.mxfp8_quantize(x, True) + assert torch.equal(values.view(torch.uint8), ref_values.view(torch.uint8)) + rows_padded = -(-m // 128) * 128 + assert scale.numel() == ref_scale.numel() == rows_padded * (k // 32) + assert torch.equal(scale, ref_scale) + + +@needs_rubin +@pytest.mark.parametrize( + "n, k, m", + [ + (4608, 8192, 8192), + (4608, 8192, 6500), + (8192, 2048, 4100), + (4096, 8192, 2048), + (8192, 4096, 1536), + (8192, 4096, 1025), + ], +) +def test_projection_matches_bf16_at_every_listed_tactic(n, k, m): + install = _install() + torch.manual_seed(0) + weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) / k**0.5 + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + copy = install.Fp8Projection(weight, min_rows=1024) + ref = F.linear(x, weight) + outputs = [copy(x, tactic) for _, tactic in copy.tactics] + for out in outputs: + assert out.shape == ref.shape and out.dtype == torch.bfloat16 + assert _cos_diff(out, ref) < 2e-3 + chosen = copy(x) + assert _cos_diff(chosen, ref) < 2e-3 + + +@needs_rubin +def test_method_dispatches_by_rows_and_captures(): + install = _install() + torch.manual_seed(0) + linear = Linear(2048, 8192, bias=False, dtype=torch.bfloat16).cuda() + with torch.no_grad(): + linear.weight.copy_(torch.randn_like(linear.weight) / 2048**0.5) + assert type(linear.quant_method) is UnquantizedLinearMethod + copy = install.Fp8Projection(linear.weight.data, min_rows=1024) + method = install.ContextFp8LinearMethod.__new__(install.ContextFp8LinearMethod) + method.__dict__.update(linear.quant_method.__dict__) + linear.quant_method = method + linear._context_fp8 = copy + assert isinstance(linear.quant_method, UnquantizedLinearMethod) + + small = torch.randn(512, 2048, device="cuda", dtype=torch.bfloat16) + # Below min_rows: the BF16 GEMM, bit for bit. + assert torch.equal(linear(small), F.linear(small, linear.weight)) + + large = torch.randn(4096, 2048, device="cuda", dtype=torch.bfloat16) + eager = linear(large) + assert torch.equal(eager, copy(large)) + assert not torch.equal(eager, F.linear(large, linear.weight)) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = linear(large) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, eager) + + +@needs_rubin +@pytest.mark.parametrize("rows", [1024, 1300, 4096]) +def test_gather_quantized_before_matches_quantized_after(rows): + """Quantizing the rows before the gather (row-major scales, swizzled after) feeds the GEMM + the same bytes as quantizing the gathered rows; below min_rows the BF16 rows pass through.""" + from tensorrt_llm.mapping import Mapping + + install = _install() + quant = importlib.import_module(f"{_PACKAGE}.quant") + torch.manual_seed(rows) + k = 8192 + x = torch.randn(rows, k, device="cuda", dtype=torch.bfloat16) + gather = install.Fp8AttentionInputGather( + Mapping(world_size=1, tp_size=1, rank=0), min_rows=1024 + ) + gathered = gather(x, rows) + assert isinstance(gathered, install.GatheredMxfp8) and tuple(gathered.shape) == (rows, k) + values, scale = quant.mxfp8_quantize(x) + assert torch.equal(gathered.values.view(torch.uint8), values.view(torch.uint8)) + assert torch.equal(gathered.scale, scale) + + weight = torch.randn(4608, k, device="cuda", dtype=torch.bfloat16) / k**0.5 + linear = Linear(k, 4608, bias=False, dtype=torch.bfloat16).cuda() + with torch.no_grad(): + linear.weight.copy_(weight) + method = install.ContextFp8LinearMethod.__new__(install.ContextFp8LinearMethod) + method.__dict__.update(linear.quant_method.__dict__) + linear.quant_method = method + linear._context_fp8 = install.Fp8Projection(linear.weight.data, min_rows=1024) + assert torch.equal(linear(gathered), linear(x)) + + small = x[:512] + passed = gather(small, 512) + assert isinstance(passed, torch.Tensor) and torch.equal(passed, small) From 1d313e3096bcad29c7054370049eacba27a29b47 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:10:59 +0000 Subject: [PATCH 07/15] [None][perf] modeling_v2 musespark_open tep: packed multi-request prep append and a parallel windowed context commit Decode steps of several requests on sm107 TEP4: - multi_request_prep_append runs the packed-geometry kernel for calls of at least packed_min_rows rows (one program per head over all sequences) instead of the per-sequence grid; the result is bitwise the same. - context_kv_commit_windowed commits the drafter's windowed context arena with one program per request and a tuned compaction for steps of at least parallel_min_requests requests; smaller steps keep the one-program commit. Both write the same bytes. - hot_replicas skips experts that were quantized at load, which have no census table. The thresholds are arguments set by the lever resolver. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../multi_request_prep_append/attention.py | 29 ++- .../replicated_experts/hot_replicas.py | 2 + .../attention/multi_request_prep_append.py | 19 +- .../ops/spec/context_kv_commit_windowed.py | 215 ++++++++++++++++-- 4 files changed, 236 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py index 9683a844e172..da7d01e8e860 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/multi_request_prep_append/attention.py @@ -81,6 +81,7 @@ from ...attention_forward import PREPROCESS_ARGS, eliding, preprocess_reachable from ...dispatch import LAUNCH_PDL from ...ops.attention.multi_request_prep_append import ( + MAX_ROWS, MAX_ROWS_PER_SEQ, MAX_SEQS, PACKED_GEOMETRY, @@ -407,13 +408,15 @@ class MultiRequestPrepAppendAttention: built-in attention's generation preprocess stand down for it (``eliding``); any other call is exactly the class behind it. Adds two plain attributes set for the duration of one forward, and two set by ``install``: the model's - ``MultiRequestCensus`` and the fewest sequences of a step it folds. + ``MultiRequestCensus`` and the fewest sequences of a step it folds, and the fewest rows + of a call on the packed kernel (``packed_prep_append``'s MIN_ROWS, None for off). """ _mr_append_metadata: Optional[object] _mr_append_folded: Optional[MultiRequestFoldedPreprocess] _mr_append_census: Optional[MultiRequestCensus] = None _mr_append_min_seqs: int = MIN_SEQS + _mr_append_packed_min_rows: Optional[int] = None def forward(self, position_ids, hidden_states, attn_metadata, *args, **kwargs): census = self._mr_append_census @@ -601,6 +604,7 @@ def _mr_append_prep( 1.0, LAUNCH_PDL, prefetch=prefetch, + packed_min_rows=self._mr_append_packed_min_rows, ) folded = MultiRequestFoldedPreprocess( qkv_out=qkv_out, @@ -632,13 +636,16 @@ def _folded_class(base: type) -> type: return type(f"{base.__name__}MultiRequestAppend", (MultiRequestPrepAppendAttention, base), {}) -def install(model: torch.nn.Module, min_seqs: int = MIN_SEQS) -> int: +def install( + model: torch.nn.Module, min_seqs: int = MIN_SEQS, packed_min_rows: Optional[int] = None +) -> int: """Re-class every layer's attention with ``MultiRequestPrepAppendAttention`` in front. Run after every other attention re-class: the mixin goes in front of whatever class each layer carries by then, and acts only on decode steps of ``min_seqs`` or more sequences (a composition whose own paths fold the append of smaller - multi-request steps passes the first count they leave to this one). + multi-request steps passes the first count they leave to this one). Calls of + ``packed_min_rows`` rows or more (``packed_prep_append``) take the packed kernel. Every layer is checked before any is re-classed. Idempotent: a model whose layers all carry the mixin is left alone. Returns the number of layers re-classed. @@ -657,6 +664,8 @@ class each layer carries by then, and acts only on decode steps of ``min_seqs`` ) attentions = [layer.self_attn for layer in model.layers] if all(isinstance(attention, MultiRequestPrepAppendAttention) for attention in attentions): + for attention in attentions: + attention._mr_append_packed_min_rows = packed_min_rows return 0 refused = [] for attention in attentions: @@ -692,13 +701,21 @@ class each layer carries by then, and acts only on decode steps of ``min_seqs`` attention._mr_append_folded = None attention._mr_append_census = census attention._mr_append_min_seqs = min_seqs - grid = ( + attention._mr_append_packed_min_rows = packed_min_rows + packed = ( "packed grid (last rows, items per program): " f"{[(last, items) for last, items, _, _ in PACKED_GEOMETRY]}, one item per program on " "o_proj-prefetching steps" - if not torch.cuda.is_available() or launch_kind(torch.device("cuda")) == "packed" - else "grid of one program per (row, head)" ) + if not torch.cuda.is_available() or launch_kind(torch.device("cuda")) == "packed": + grid = packed + elif launch_kind(torch.device("cuda"), MAX_ROWS, packed_min_rows) == "packed": + grid = ( + f"grid of one program per (row, head) up to {packed_min_rows - 1} rows, " + f"above it the {packed}" + ) + else: + grid = "grid of one program per (row, head)" logger.info( f"[musespark_open] multi_request_prep_append: the K/V append, E4M3 query and FMHA scales " f"of {len(attentions)} layers run in the decode prep on steps of {min_seqs}..{MAX_SEQS} sequences " diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/hot_replicas.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/hot_replicas.py index fcd0384b8f54..afd59d6fb844 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/hot_replicas.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/replicated_experts/hot_replicas.py @@ -145,6 +145,8 @@ def eligible(moe, mapping: Mapping) -> Optional[str]: return reason if _ring_buddy.expert_format(moe) != "nvfp4": return "the experts are not the native NVFP4 leaf" + if getattr(moe, "experts_quantized_at_load", False): + return "the experts were quantized at load; the census covers the checkpoint's NVFP4 layers" return None diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py index e9fd45808a1f..dc122dc6c792 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/attention/multi_request_prep_append.py @@ -63,7 +63,9 @@ one-program-per-(row, head) grid element for element (the same loads, the same additions in the same order, the same rounding points and stores), so the bytes written do not depend on the geometry. On sm107 the grid is one program per -(row, head), one warp each (``_multi_request_prep_append_grid_kernel``). +(row, head), one warp each (``_multi_request_prep_append_grid_kernel``); a call of +``packed_min_rows`` rows or more (``packed_prep_append``'s MIN_ROWS) takes the +packed kernel above on either. """ from __future__ import annotations @@ -411,8 +413,10 @@ def _sm_of(index: int) -> int: return 10 * major + minor -def launch_kind(device: torch.device) -> str: - """``LAUNCH_KIND_BY_SM`` for ``device``. +def launch_kind(device: torch.device, rows: int = 0, packed_min_rows: Optional[int] = None) -> str: + """``LAUNCH_KIND_BY_SM`` for ``device``, or ``"packed"`` for a call of at least + ``packed_min_rows`` rows (None: never). The packed kernel on ``PACKED_GEOMETRY`` writes + the grid kernel's bytes. Raises: ValueError: on a device ``LAUNCH_KIND_BY_SM`` does not list. @@ -424,6 +428,8 @@ def launch_kind(device: torch.device) -> str: f"multi_request_prep_append: no launch for sm{sm} " f"(supported: {sorted(LAUNCH_KIND_BY_SM)})" ) + if packed_min_rows is not None and rows >= packed_min_rows: + return "packed" return kind @@ -661,7 +667,7 @@ class MultiRequestPrepAppend(OpWrapper): the target then skips. Grid on sm103: the rows * (H + G) (row, head) items in row-major order, PACKED_GEOMETRY's items, warps and lanes per program by row count (with a prefetch, one program per item); on sm107: (rows, H + G), one warp per - program. + program, up to 32 rows, and the sm103 grid above. Bit-identical, gated by its own tests: qkv_out and gate_out against preattn_prep; the cache bytes, the E4M3 query bytes, both scale words and @@ -708,6 +714,7 @@ def __call__( o_scale: float = 1.0, enable_pdl: bool = False, prefetch: Optional[torch.Tensor] = None, + packed_min_rows: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: rows, width = qkv.shape page = kv_pool.shape[2] if kv_pool.dim() == 4 else 0 @@ -770,7 +777,7 @@ def __call__( raise ValueError("multi_request_prep_append: position_ids without a cos/sin table") heads = num_q_heads + num_kv_heads num_items = rows * heads - kind = launch_kind(qkv.device) + kind = launch_kind(qkv.device, rows, packed_min_rows) pf_base, pf_bytes, pf_chunk = 0, 0, 16 programs_at_least = 0 if prefetch is not None: @@ -916,6 +923,7 @@ def reference( o_scale: float = 1.0, enable_pdl: bool = False, prefetch: Optional[torch.Tensor] = None, + packed_min_rows: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """``preattn_prep.reference`` (float64 between its rounding points), then the same in-place writes from its BF16 values. @@ -997,6 +1005,7 @@ def is_valid( o_scale: float = 1.0, enable_pdl: bool = False, prefetch: Optional[torch.Tensor] = None, + packed_min_rows: Optional[int] = None, ) -> None: """Inputs the kernel takes and answers wrongly, with no error. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit_windowed.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit_windowed.py index 91207a8f03b1..47f1b9eb5713 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit_windowed.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit_windowed.py @@ -20,7 +20,9 @@ stream that does not wait for the verify. This entry is what is left once the accepted counts exist, as two kernels: -* the commit (one program, like ``context_kv_commit``'s): the zeroing of the +* the commit (one program, like ``context_kv_commit``'s, or with + ``parallel_min_requests`` a parallel zeroing and a vectorized length launch on + calls of at least that many requests): the zeroing of the masked rows in place at their physical columns, the length update, each request's shift, the base update and the new physical lengths; * the compaction (one program per request, layer, arena and chunk of a @@ -38,6 +40,8 @@ from __future__ import annotations +from typing import Optional + import torch import triton import triton.language as tl @@ -53,6 +57,11 @@ #: head_dim`` elements split into chunks of this many. COMPACT_CHUNK = 16 +#: The compaction's (chunk, columns per tile, warps) behind the parallel commit +#: (``parallel_min_requests``); the one-program commit's compaction is +#: (``COMPACT_CHUNK``, ``COMPACT_COLUMNS``, 4). Both write the same bytes. +COMPACT_TUNED = (64, 256, 8) + @triton.jit def _commit_windowed_kernel( @@ -135,6 +144,101 @@ def _commit_windowed_kernel( tl.store(num_ctx_ptr + g, tl.load(ctx_len_ptr + slot) - tl.load(ctx_base_ptr + slot)) +@triton.jit +def _commit_zero_kernel( + k_ptr, + v_ptr, + ctx_len_ptr, + ctx_base_ptr, + slots_ptr, + accepted_ptr, + new_len_ptr, + col_limit, + k_slot_stride, + k_layer_stride, + k_col_stride, + v_slot_stride, + v_layer_stride, + v_col_stride, + ROWS: tl.constexpr, + LAYERS: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, + PAIRS: tl.constexpr, + HAS_V: tl.constexpr, +): + """The zeroing of ``_commit_windowed_kernel`` for request ``program_id(0)``, and its + new length ``ctx_len[slot] + accepted`` into ``new_len`` (lengths and bases only read).""" + g = tl.program_id(0) + slot = tl.load(slots_ptr + g).to(tl.int64) + ctx = tl.load(ctx_len_ptr + slot) + phys = ctx - tl.load(ctx_base_ptr + slot) + accepted = tl.load(accepted_ptr + g).to(tl.int64) + for start in tl.static_range(0, ROWS * LAYERS, PAIRS): + k_cols, zeroed = _columns( + k_ptr, slot, phys, accepted, col_limit - 1, k_slot_stride, k_layer_stride, + k_col_stride, start, ROWS, LAYERS, WIDTH, BLOCK, PAIRS, + ) # fmt: skip + if HAS_V: + v_cols, _ = _columns( + v_ptr, slot, phys, accepted, col_limit - 1, v_slot_stride, v_layer_stride, + v_col_stride, start, ROWS, LAYERS, WIDTH, BLOCK, PAIRS, + ) # fmt: skip + k = tl.load(k_cols, mask=zeroed, other=0.0) + v = tl.load(v_cols, mask=zeroed, other=0.0) + tl.store(k_cols, _times_zero(k), mask=zeroed) + tl.store(v_cols, _times_zero(v), mask=zeroed) + else: + k = tl.load(k_cols, mask=zeroed, other=0.0) + tl.store(k_cols, _times_zero(k), mask=zeroed) + tl.store(new_len_ptr + g, ctx + accepted) + + +@triton.jit +def _commit_lengths_kernel( + ctx_len_ptr, + ctx_base_ptr, + slots_ptr, + new_len_ptr, + new_base_ptr, + shift_ptr, + num_ctx_ptr, + num_gens, + num_slots, + max_ctx, + keep, + GENS: tl.constexpr, + SLOTS: tl.constexpr, +): + """The length, shift and base updates of ``_commit_windowed_kernel``, vectorized over the + requests: a repeated slot (CUDA-graph padding) is written by its last request only, as + the request-order scatter leaves it.""" + gi = tl.arange(0, GENS) + gm = gi < num_gens + slot = tl.load(slots_ptr + gi, mask=gm, other=-1).to(tl.int64) + later = (gi[None, :] > gi[:, None]) & gm[None, :] & (slot[None, :] == slot[:, None]) + writer = gm & (tl.sum(later.to(tl.int32), 1) == 0) + tl.store(ctx_len_ptr + slot, tl.load(new_len_ptr + gi, mask=gm, other=0), mask=writer) + tl.debug_barrier() + si = tl.arange(0, SLOTS) + sm = si < num_slots + lengths = tl.load(ctx_len_ptr + si, mask=sm, other=0) + tl.store(ctx_len_ptr + si, tl.minimum(lengths, max_ctx), mask=sm) + tl.debug_barrier() + length = tl.load(ctx_len_ptr + slot, mask=gm, other=0) + base = tl.load(ctx_base_ptr + slot, mask=gm, other=0) + shift = tl.where(keep > 0, tl.maximum(length - base - keep, 0), 0) + tl.store(shift_ptr + gi, shift, mask=gm) + tl.store(new_base_ptr + gi, base + shift, mask=gm) + tl.debug_barrier() + tl.store(ctx_base_ptr + slot, base + shift, mask=writer) + tl.debug_barrier() + num_ctx = tl.load(ctx_len_ptr + slot, mask=gm, other=0) - tl.load( + ctx_base_ptr + slot, mask=gm, other=0 + ) + tl.store(num_ctx_ptr + gi, num_ctx, mask=gm) + + @triton.jit def _compact_kernel( k_ptr, @@ -223,7 +327,10 @@ class ContextKvCommitWindowed(OpWrapper): ``COMPACT_CHUNK`` elements); a program loops over the slot's columns in tiles of ``COMPACT_COLUMNS`` in increasing order, loading a tile, a CTA barrier, then storing it; it reads its request's shift from the commit's - output, so its launch follows the commit on the stream. + output, so its launch follows the commit on the stream. With + ``parallel_min_requests`` (``windowed_commit_parallel``'s MIN_REQUESTS) a call of + at least that many requests runs the commit as a parallel zeroing launch plus a + vectorized length launch, and the compaction at ``COMPACT_TUNED``; the same bytes. Call site: ``context_kv_fork.ContextKvFork.commit`` under the ``context_kv_fork`` fast path when the worker's arena is windowed @@ -259,15 +366,60 @@ def __call__( max_ctx: int, col_limit: int, keep: int, + parallel_min_requests: Optional[int] = None, ) -> None: num_gens = slots.numel() if num_gens == 0: return + parallel = parallel_min_requests is not None and num_gens >= parallel_min_requests layers, width = k_arena.shape[1], k_arena.shape[3] * k_arena.shape[4] scratch = torch.empty(3, num_gens, dtype=torch.int64, device=ctx_len.device) has_v = v_arena is not None values = v_arena if has_v else k_arena - _commit_windowed_kernel[(1,)]( + if parallel: + self._commit_parallel( + k_arena, values, has_v, ctx_len, ctx_base, slots, num_accepted, num_ctx_out, + scratch, rows, max_ctx, col_limit, keep, layers, width, + ) # fmt: skip + else: + self._commit_one_program( + k_arena, values, has_v, ctx_len, ctx_base, slots, num_accepted, num_ctx_out, + scratch, rows, max_ctx, col_limit, keep, layers, width, + ) # fmt: skip + if keep <= 0: + return + if parallel: + chunk, tile, warps = COMPACT_TUNED + else: + chunk, tile, warps = COMPACT_CHUNK, COMPACT_COLUMNS, 4 + chunk = min(chunk, triton.next_power_of_2(width)) + _compact_kernel[(num_gens, layers * (2 if has_v else 1), triton.cdiv(width, chunk))]( + k_arena, + values, + slots, + scratch[2], + k_arena.shape[2], + k_arena.stride(0), + k_arena.stride(1), + k_arena.stride(2), + values.stride(0), + values.stride(1), + values.stride(2), + LAYERS=layers, + WIDTH=width, + CHUNK=chunk, + TILE=tile, + GENS=triton.next_power_of_2(max(num_gens, 2)), + HAS_V=has_v, + num_warps=warps, + ) + + def _commit_parallel( + self, k_arena, values, has_v, ctx_len, ctx_base, slots, num_accepted, num_ctx_out, + scratch, rows, max_ctx, col_limit, keep, layers, width, + ) -> None: # fmt: skip + num_gens = slots.numel() + _commit_zero_kernel[(num_gens,)]( k_arena, values, ctx_len, @@ -275,14 +427,7 @@ def __call__( slots, num_accepted, scratch[0], - scratch[1], - scratch[2], - num_ctx_out, - num_gens, - ctx_len.numel(), - max_ctx, col_limit, - keep, k_arena.stride(0), k_arena.stride(1), k_arena.stride(2), @@ -297,26 +442,55 @@ def __call__( HAS_V=has_v, num_warps=4, ) - if keep <= 0: - return - chunk = min(COMPACT_CHUNK, triton.next_power_of_2(width)) - _compact_kernel[(num_gens, layers * (2 if has_v else 1), triton.cdiv(width, chunk))]( + _commit_lengths_kernel[(1,)]( + ctx_len, + ctx_base, + slots, + scratch[0], + scratch[1], + scratch[2], + num_ctx_out, + num_gens, + ctx_len.numel(), + max_ctx, + keep, + GENS=triton.next_power_of_2(max(num_gens, 2)), + SLOTS=triton.next_power_of_2(max(ctx_len.numel(), 2)), + num_warps=4, + ) + + def _commit_one_program( + self, k_arena, values, has_v, ctx_len, ctx_base, slots, num_accepted, num_ctx_out, + scratch, rows, max_ctx, col_limit, keep, layers, width, + ) -> None: # fmt: skip + num_gens = slots.numel() + _commit_windowed_kernel[(1,)]( k_arena, values, + ctx_len, + ctx_base, slots, + num_accepted, + scratch[0], + scratch[1], scratch[2], - k_arena.shape[2], + num_ctx_out, + num_gens, + ctx_len.numel(), + max_ctx, + col_limit, + keep, k_arena.stride(0), k_arena.stride(1), k_arena.stride(2), values.stride(0), values.stride(1), values.stride(2), + ROWS=rows, LAYERS=layers, WIDTH=width, - CHUNK=chunk, - TILE=COMPACT_COLUMNS, - GENS=triton.next_power_of_2(max(num_gens, 2)), + BLOCK=triton.next_power_of_2(width), + PAIRS=tile_pairs(rows, layers, width), HAS_V=has_v, num_warps=4, ) @@ -334,6 +508,7 @@ def reference( max_ctx: int, col_limit: int, keep: int, + parallel_min_requests: Optional[int] = None, ) -> None: """The built-in torch ops, in place: mask the stored rows, update the lengths, compact. @@ -381,7 +556,11 @@ def is_valid( max_ctx: int, col_limit: int, keep: int, + parallel_min_requests: Optional[int] = None, ) -> None: + assert parallel_min_requests is None or parallel_min_requests >= 1, ( + "parallel_min_requests must be None or at least 1" + ) for name, arena in (("k_arena", k_arena), ("v_arena", v_arena)): if arena is None and name == "v_arena": continue From a823b4352d8193a1a7c147e6a584b051ae4b4a01 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:11:01 +0000 Subject: [PATCH 08/15] [None][perf] modeling_v2 musespark_open tep: wide direct transports for decode steps of 16 to 32 rows Decode steps of two to four requests (16 to 32 verify rows) on sm107 TEP4: - wide_prologue routes each row with its own one-request bits on the wide router (router shard rows), so a row routes as it would alone; - pool_into_route and unpool_into_boundary run the direct pool and unpool transports with the half expert window at 16, 24 and 32 rows; - attention_output_into_hop publishes o_proj directly into the wide post-attention hop transport at 16, 24 and 32 rows; - finalize_transport serves finalize calls of 17 to 32 tokens. The 8-row layout shares the kernel builders; its arithmetic is unchanged. Which of these paths a deployment serves is decided by the lever resolver. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../kernels/finalize_transport/transport.py | 35 +- .../kernels/moe_prologue/prologue.py | 62 ++- .../kernels/pool_into_route/route.py | 180 +++++++- .../kernels/unpool_into_boundary/boundary.py | 81 +++- .../kernels/wide_prologue/prologue.py | 33 +- .../ops/allreduce/post_attn_hop.py | 226 ++++++++-- .../ops/comm/pool_direct_route.py | 406 +++++++++++++++--- .../ops/comm/unpool_direct_boundary.py | 143 +++++- .../ops/moe/finalize_allreduce_transport.py | 260 +++++++++-- .../ops/moe/router_shard.py | 251 ++++++++++- 10 files changed, 1455 insertions(+), 222 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_transport/transport.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_transport/transport.py index bca71b52b6f3..845f7afd9f9b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_transport/transport.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/finalize_transport/transport.py @@ -16,8 +16,13 @@ compilation on every rank) and puts it in every layer's attribute. Its ``finalize`` has the built-in signature: calls of 1..16 tokens run ``ops/moe/finalize_allreduce_transport``, bit for bit the built-in kernel's -result; larger calls run the built-in kernel. No module is re-classed, so the -fast paths that re-class the MoE (``projection_gathers``) or patch its experts (``replicated_experts``) +result; larger calls run the built-in kernel. With ``wide_finalize_transport`` +(``install(model, wide=True)``, TP4 only) the workspace also builds a second +one laid out for 32 tokens, so the calls of 17..32 tokens (decode steps of +three or four requests at K = 7) run the same kernel there, bit for bit the +built-in TP4 kernel too, and only larger calls run the built-in kernel. No +module is re-classed, so the fast paths that re-class the MoE +(``projection_gathers``) or patch its experts (``replicated_experts``) compose with it in any order. """ @@ -37,8 +42,11 @@ FEATURES = {"finalize_transport": _fin.DEFAULT} -def refusal(model: MuseSparkOpenModel) -> Optional[str]: - """Why ``install`` cannot replace this model's fused finalize, or None. Changes nothing.""" +def refusal(model: MuseSparkOpenModel, wide: bool = False) -> Optional[str]: + """Why ``install`` cannot replace this model's fused finalize, or None. Changes nothing. + + ``wide`` asks for the 17..32-token calls as well, which only a TP4 group has. + """ legacy = getattr(model, ATTRIBUTE, None) if legacy is None: return ( @@ -55,17 +63,21 @@ def refusal(model: MuseSparkOpenModel) -> Optional[str]: return f"layers {others} do not call the model's fused finalize workspace" if not _fin.HAS_CUTE_DSL: return "the CuTe DSL or the built-in finalize module is unavailable" + if wide: + tp_size = getattr(model.layers[0].mapping, "tp_size", None) + if tp_size not in _fin.LAYOUT_RANK_COUNTS[_fin.WIDE_MAX_TOKENS]: + return f"the {_fin.WIDE_MAX_TOKENS}-token layout serves TP4, got TP{tp_size}" return None -def agreed_refusal(model: MuseSparkOpenModel) -> Optional[str]: +def agreed_refusal(model: MuseSparkOpenModel, wide: bool = False) -> Optional[str]: """Why ``install`` refuses, as every TP rank agrees, or None (collective over the TP group). Each rank's ``refusal`` is gathered, so that a reason only some ranks see (ranks running different images or environments) makes every rank refuse before the workspace's collective construction. Changes nothing. """ - reason = refusal(model) + reason = refusal(model, wide) local = None if reason is None else f"[musespark_open] finalize_transport cannot run: {reason}" everyone = Distributed.get(model.layers[0].mapping).tp_allgather(local, small_payload=True) refused = [(rank, why) for rank, why in enumerate(everyone) if why is not None] @@ -78,22 +90,27 @@ def agreed_refusal(model: MuseSparkOpenModel) -> Optional[str]: return f"{'; '.join(reasons)} (on TP ranks {ranks} of {len(everyone)})" -def install(model: MuseSparkOpenModel): +def install(model: MuseSparkOpenModel, wide: bool = False): """Put a ``FinalizeTransportWorkspace`` in every layer's finalize attribute (collective). + With ``wide`` the workspace also serves the calls of 17..32 tokens on its + second, 32-token workspace (``FinalizeTransportWorkspace.serve_wide_calls``). + Raises: ValueError: if ``refusal`` names a reason on some rank (on every rank alike, from ``agreed_refusal``, before the workspace is built). RuntimeError: if the workspace could not be built on some rank (on every rank). """ - reason = agreed_refusal(model) + reason = agreed_refusal(model, wide) if reason is not None: raise ValueError(reason) legacy = getattr(model, ATTRIBUTE) transport = FEATURES["finalize_transport"] workspace = _fin.FinalizeTransportWorkspace(legacy, legacy._mapping, transport) + if wide: + workspace.serve_wide_calls(legacy._mapping) for layer in model.layers: setattr(layer.feed_forward, ATTRIBUTE, workspace) - _fin.log_once_enabled(len(model.layers), transport) + _fin.log_once_enabled(len(model.layers), transport, wide) return workspace diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/moe_prologue/prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/moe_prologue/prologue.py index 8b3625ba30cf..52ec24414334 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/moe_prologue/prologue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/moe_prologue/prologue.py @@ -134,6 +134,13 @@ class MuseSparkOpenTepPrologueMoE(MuseSparkOpenTepGatherMoE): #: ``wide_prologue``: the TP4 wide inbox (``WIDE_MAX_ROWS`` rows) of the 16..32-row #: prologue, or None (those steps keep the built-in prepare). _wide_inbox: Optional[ProjectionInbox] = None + #: ``wide_prologue``: whether the 16..32-row router shard gives each row its + #: one-request bits (``router_shard_rows``) instead of the built-in tensor-core + #: router's (``router_shard_wide``); the ``wide_router_per_row`` lever. + _wide_router_rows: bool = False + #: ``wide_prologue``: whether ``pool_into_route`` and ``unpool_into_boundary`` also + #: build their 16..32-row direct transports; the ``wide_direct_transports`` lever. + _wide_direct: bool = False def prepare_moe_inputs(self, hidden_states: torch.Tensor): if self._wide_inbox is not None: @@ -221,13 +228,39 @@ def prepare_moe_inputs(self, hidden_states: torch.Tensor): self._prologue_record_routing(out, place, state, world) return out.logits, Fp4QuantizedTensor(out.fp4, out.scale, is_sf_swizzled=False) + def _wide_router_shard( + self, hidden_states: torch.Tensor, rank: int, warps: int, out: torch.Tensor, world: int + ) -> None: + """This rank's router logit columns of a 16..32-row call into ``out`` (``[rows, L]`` FP32). + + By default ``router_shard_wide``: the built-in 16..32-row router's bits (its TF32x3 + tensor-core split-K and reduce). With ``_wide_router_rows``, ``router_shard_rows`` at + ``warps``: each row's logits are the built-in small-batch router's for that row in a + one-request call, one launch and no reduce, but other bits than the built-in + 16..32-row launch's. + """ + if self._wide_router_rows: + _shard.router_shard_rows( + hidden_states, + self.router.weight, + rank, + self.metap_multiplier, + warps, + out, + ranks=world, + ) + else: + _shard.router_shard_wide( + hidden_states, self.router.weight, rank, self.metap_multiplier, out, ranks=world + ) + def _wide_prepare(self, hidden_states: torch.Tensor): """``wide_prologue``: the prologue of a 16..32-row decode step, or None (nothing launched). Where the built-in decode would run its tensor-core router and the fused routing on a pool-deferring call of 16..32 rows: this rank's router columns - (``router_shard_wide``, the built-in tensor-core bits) on the router fork's side - stream, the pool slice GEMM on the current one, and ``pool_gather_route`` on the + (``_wide_router_shard``) on the router fork's side stream, the pool slice GEMM on the + current one, and ``pool_gather_route`` on the wide inbox -- the pool gather + quantization, the logit exchange and the routing (no placement: above 8 rows the experts are the built-in ones and read only the routing memo). The pool slice is this rank's columns of the built-in replicated @@ -246,6 +279,13 @@ def _wide_prepare(self, hidden_states: torch.Tensor): if args is None: return None gain, global_scale = args + # One request's warp count (the built-in launch's at 8 rows), for router_shard_rows. + warps = _shard.builtin_warps( + _shard.MAX_ROWS, getattr(self, "_prepare_stream", None) is not None + ) + if self._pool_direct_armed is not None: + # pool_into_route's wide site: the same chain with the producers publishing. + return self._pool_direct_armed.prepare(self, hidden_states, warps, gain, global_scale) routing = self.experts.routing_method side = self._router_stream local = _shard.local_experts(world) @@ -255,25 +295,11 @@ def _wide_prepare(self, hidden_states: torch.Tensor): (rows, local), device=hidden_states.device, dtype=torch.float32 ) if side is None: - _shard.router_shard_wide( - hidden_states, - self.router.weight, - inbox.rank, - self.metap_multiplier, - logits_local, - ranks=world, - ) + self._wide_router_shard(hidden_states, inbox.rank, warps, logits_local, world) else: side.wait_stream(main) with torch.cuda.stream(side): - _shard.router_shard_wide( - hidden_states, - self.router.weight, - inbox.rank, - self.metap_multiplier, - logits_local, - ranks=world, - ) + self._wide_router_shard(hidden_states, inbox.rank, warps, logits_local, world) with self.nvtx("pool"): projected = self._replicated_projection( self.pooler, hidden_states, self._pool_gather_workspace diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/pool_into_route/route.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/pool_into_route/route.py index 7e34e85dd923..0756cc6c1118 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/pool_into_route/route.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/pool_into_route/route.py @@ -51,9 +51,22 @@ ``moe_prologue`` then runs ``DirectPoolSite.prepare`` in place of the tagged chain. Everything is disarmed on any exit. +Steps of two to four requests (``wide_prologue``, TP4): when every layer holds the wide +route inbox and the wide direct transports (the ``wide_direct_transports`` lever), the +group also gets a second transport on the wide layout +(``PoolDirectWorkspace(max_rows=WIDE_MAX_ROWS)``) with its own sites and buffer +sequence, and a 16..32-row forward captured into a CUDA graph arms those. The prologue's +wide path (``_wide_prepare``) then runs ``DirectPoolSite.prepare`` with the wide site: +the prologue's wide router shard (``_wide_router_shard``) writes its ``[rows, L]`` logits +into the record's logit block on the side stream, the pool GEMM (``wide_dense_gemms``' split-K) writes the slot, and the +routed kernel runs the wide tagged kernel's statements. On a half-window layer +(``half_window_experts``) the routed kernel also places: its CTAs merge their rows' census +in rank-local memory and each writes its rows' local ids, so the step's placement is +recorded here and ``ring_buddy_assign(half=True)`` does not run. Same bits. + An armed layer whose call does not take the prologue's routed path raises: its peers would poll for a record nobody publishes. Eager steps and other row counts keep -the tagged chain; the two transports never share memory. +the tagged chain; the transports never share memory. """ from __future__ import annotations @@ -69,6 +82,7 @@ from ... import dispatch as _tp_agreement from ...ops.comm import pool_direct_route as _pdr from ...ops.comm import unpool_direct_boundary as _direct +from ...ops.comm.projection_inbox import WIDE_MAX_ROWS from ...ops.moe import router_shard as _shard from ..dense_gemms.dispatcher import ( FLASHINFER_FAMILIES, @@ -105,15 +119,22 @@ def __init__( self.dispatcher = dispatcher def prepare(self, moe, hidden_states: torch.Tensor, warps: int, gain, global_scale): - """The prologue MoE's ``prepare_moe_inputs`` result for an armed 8-row call, through the transport. + """The prologue MoE's ``prepare_moe_inputs`` result for an armed call, through the transport. ``moe`` is the layer's prologue MoE with the router fork's side stream. ``warps``, ``gain`` and ``global_scale`` are what its routed path computed for this - call. + call (its 8-row path, or its wide path on a wide site). """ from ..moe_prologue.prologue import _placement rows = hidden_states.shape[0] + if self.workspace.wide: + if not _shard.WIDE_MIN_ROWS <= rows <= self.workspace.max_rows: + raise RuntimeError( + f"[musespark_open] pool_into_route: a layer armed for wide calls got {rows} rows; " + "its peers poll the transport for its records" + ) + return self._prepare_wide(moe, hidden_states, warps, gain, global_scale) if rows != ROWS: raise RuntimeError( f"[musespark_open] pool_into_route: a layer armed for {ROWS}-row calls got {rows} rows; its peers " @@ -174,35 +195,98 @@ def prepare(self, moe, hidden_states: torch.Tensor, warps: int, gain, global_sca moe._prologue_record_routing(out, place, state) return out.logits, Fp4QuantizedTensor(out.fp4, out.scale, is_sf_swizzled=False) + def _projected_into_slot(self, moe, hidden_states: torch.Tensor, slot: torch.Tensor) -> None: + self.dispatcher.publish = slot + try: + projected = moe._replicated_projection( + moe.pooler, hidden_states, moe._pool_gather_workspace + ) + finally: + self.dispatcher.publish = None + if projected.data_ptr() != slot.data_ptr(): + raise RuntimeError( + "[musespark_open] pool_into_route: the pool GEMM of an armed layer did not write the transport " + "slot (it bypassed the dense_gemms pool dispatcher); its peers poll for that slot" + ) + + def _prepare_wide(self, moe, hidden_states: torch.Tensor, warps: int, gain, global_scale): + """``prepare`` of a 16..32-row call on the wide layout: ``_wide_prepare``'s chain with + the producers publishing, and on a half-window layer the placement too.""" + rows = hidden_states.shape[0] + place, static_from, state = _wide_placement(moe, self.workspace.world, rows) + workspace, buffer = self.workspace, self.buffer + routing = moe.experts.routing_method + side = moe._router_stream + main = torch.cuda.current_stream() + with moe.nvtx("router"): + block = workspace.block(buffer, rows) + side.wait_stream(main) + with torch.cuda.stream(side): + moe._wide_router_shard(hidden_states, workspace.rank, warps, block, workspace.world) + with moe.nvtx("pool"): + self._projected_into_slot(moe, hidden_states, workspace.slot(buffer, rows)) + main.wait_stream(side) + bias = routing.e_score_correction_bias + _require_aligned(gain=gain, bias=bias) + out = _pdr.pool_direct_route( + workspace, + buffer, + rows, + gain, + global_scale, + moe.pre_expert_norm.variance_epsilon, + bias, + routing.renorm_eps, + place, + static_from, + enable_pdl=moe._prologue_enable_pdl, + return_gathered=False, + ) + moe._prologue_record_routing(out, place, state, workspace.world) + return out.logits, Fp4QuantizedTensor(out.fp4, out.scale, is_sf_swizzled=False) + class DirectPool: - """What ``DecodeTarget`` holds under ``pool_into_route``: the workspace and, once loaded, its sites. + """What ``DecodeTarget`` holds under ``pool_into_route``: the workspaces and, once loaded, their sites. - ``sites`` is ``(moe, site)`` per layer that runs the routed path, in call order. + ``sites`` is ``(moe, site)`` per layer that runs the routed path, in call order; + ``wide_sites`` the same layers on the wide workspace (``wide_workspace``, present when + every layer holds the wide route inbox), else empty. """ - def __init__(self, workspace: _pdr.PoolDirectWorkspace) -> None: + def __init__( + self, + workspace: _pdr.PoolDirectWorkspace, + wide_workspace: Optional[_pdr.PoolDirectWorkspace] = None, + ) -> None: self.workspace = workspace + self.wide_workspace = wide_workspace self.sites: Tuple[Tuple[torch.nn.Module, DirectPoolSite], ...] = () + self.wide_sites: Tuple[Tuple[torch.nn.Module, DirectPoolSite], ...] = () - @staticmethod - def takes(rows: int) -> bool: - """Whether a decode forward of ``rows`` rows runs on the direct transport. + def takes(self, rows: int) -> bool: + """Whether a decode forward of ``rows`` rows runs on a direct transport. - Exactly 8 rows, inside a CUDA graph capture. Every rank captures the same - graphs with the same row counts, so either every rank arms or none does. + Inside a CUDA graph capture: exactly 8 rows, or 16..32 rows when the wide sites + are installed. Every rank captures the same graphs with the same row counts and + installs the same sites, so either every rank arms or none does. """ - return rows == ROWS and torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + return False + if rows == ROWS: + return True + return bool(self.wide_sites) and _shard.WIDE_MIN_ROWS <= rows <= WIDE_MAX_ROWS @contextlib.contextmanager def armed(self, rows: int): - """Arm every site for one forward of ``rows`` rows; disarm on any exit.""" - for moe, site in self.sites: + """Arm every site of the ``rows``-row transport for one forward; disarm on any exit.""" + sites = self.sites if rows <= ROWS else self.wide_sites + for moe, site in sites: moe._pool_direct_armed = site try: yield finally: - for moe, _ in self.sites: + for moe, _ in sites: moe._pool_direct_armed = None def variants(self, layers: Sequence) -> list: @@ -278,25 +362,69 @@ def install_sites(self, model, mapping) -> int: ) if reason is not None: raise ValueError(reason) from failed + wide = self.wide_workspace + wide_variants = [] + if wide is not None: + wide_variants = self.wide_variants(routed) + failed = None + try: + wide.build_variants(wide_variants) + except Exception as error: # noqa: BLE001 - every failure must reach the agreement + failed = error + reason = _tp_agreement.agreed_refusal( + mapping, + None + if failed is None + else f"[musespark_open] pool_into_route could not build its wide routed kernel: {failed!r}", + ) + if reason is not None: + raise ValueError(reason) from failed buffers = _direct.buffer_sequence(len(routed)) - sites = [] + sites, wide_sites = [], [] for layer, buffer in zip(routed, buffers): moe = layer.feed_forward dispatcher = _pool_dispatcher(moe) claim_publish(dispatcher, PUBLISH_POLICY) moe._pool_direct_armed = None sites.append((moe, DirectPoolSite(self.workspace, buffer, dispatcher))) + if wide is not None: + wide_sites.append((moe, DirectPoolSite(wide, buffer, dispatcher))) self.sites = tuple(sites) + self.wide_sites = tuple(wide_sites) self.workspace.start_guard_watch() + if wide is not None: + wide.start_guard_watch() logger.info( f"[musespark_open] pool_into_route: the pool GEMM and the forked router shard publish into the routed pool " f"kernel's direct transport on {len(sites)} of {len(layers)} layers (8-row decode " - f"graphs; buffers alternate per layer, the last on buffer {buffers[-1]}; " - f"{len(variants)} kernel variants); eager steps and other row counts keep the tagged " - "chain" + f"graphs{'' if wide is None else f', and 16..{WIDE_MAX_ROWS}-row ones on the wide layout'}; " + f"buffers alternate per layer, the last on buffer {buffers[-1]}; " + f"{len(variants) + len(wide_variants)} kernel variants); eager steps and other row " + "counts keep the tagged chain" ) return len(sites) + def wide_variants(self, layers: Sequence) -> list: + """The wide routed kernel's variants the layers' wide calls launch (``_wide_placement``).""" + keys = set() + for layer in layers: + moe = layer.feed_forward + for rows in (_shard.WIDE_MIN_ROWS, WIDE_MAX_ROWS): + place, static_from, _ = _wide_placement(moe, self.wide_workspace.world, rows) + keys.add(_pdr.variant(place, static_from, bool(moe._prologue_enable_pdl), False)) + return sorted(keys) + + +def _wide_placement(moe, world: int, rows: int): + """``(kind, static_from, state)`` of a wide call: the half window's placement on a layer + whose placement serves ``rows`` rows, else routing only (``ring_buddy_assign`` places).""" + from ..moe_prologue.prologue import _placement + + place, static_from, state = _placement(moe, world) + if place == _pdr.PLACE_HALF and state is not None and rows <= state.max_rows: + return place, static_from, state + return _pdr.PLACE_NONE, _pdr.LOCAL_EXPERTS, None + def _pool_dispatcher(moe) -> Optional[DenseGemmDispatcher]: pooler = getattr(moe, "pooler", None) @@ -386,4 +514,16 @@ def install(model, plan, mapping=None) -> DirectPool: reason = _tp_agreement.agreed_refusal(mapping, reason) if reason is not None: raise ValueError(reason) - return DirectPool(_pdr.shared_workspace(layers[0].mapping)) + group = layers[0].mapping + workspace = _pdr.shared_workspace(group) + # The wide layout when wide_prologue gave every layer its wide route inbox and the + # wide_direct_transports lever (TP4; construction-time state agreed over the TP + # group by wide_prologue.install, the same on every rank). + wide = None + if all( + getattr(layer.feed_forward, "_wide_inbox", None) is not None + and getattr(layer.feed_forward, "_wide_direct", False) + for layer in layers + ): + wide = _pdr.shared_workspace(group, max_rows=WIDE_MAX_ROWS) + return DirectPool(workspace, wide) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/unpool_into_boundary/boundary.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/unpool_into_boundary/boundary.py index 32856bd37cca..284e8f8e578e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/unpool_into_boundary/boundary.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/unpool_into_boundary/boundary.py @@ -34,8 +34,15 @@ last fused layer takes buffer 2, so the next step's first call never lands on the buffer the previous call used. The reuse order then rests on the transport's own calls, not on the collectives between the steps' boundaries. +* Steps of two to four requests (``wide_prologue``, TP4): when every fused join holds + the wide unpool inbox and its layer the wide direct transports (the + ``wide_direct_transports`` lever), the group also gets a transport on the wide layout + (``UnpoolDirectWorkspace(max_rows=WIDE_MAX_ROWS)``, its own buffer sequence over the + same layers), and a captured forward of 16, 24 or 32 rows arms that one: the unpool + GEMM (``wide_dense_gemms``' split-K) writes the layer's wide slot and the join runs the + boundary at that row count. * Eager steps, other row counts and the last layer (no fused boundary) keep - the projection gathers' tagged inbox and path; the two transports never share memory. + the projection gathers' tagged inbox and path; the transports never share memory. The join raises instead of falling back when it is armed and the call is not the fused boundary's: its GEMM has already published into every rank's slot, @@ -52,6 +59,7 @@ from ... import dispatch as _tp_agreement from ...ops.comm import unpool_direct_boundary as _direct from ...ops.comm import unpool_gather_boundary as _unpool +from ...ops.comm.projection_inbox import WIDE_MAX_ROWS from ..dense_gemms.dispatcher import ( FLASHINFER_FAMILIES, DenseGemmDispatcher, @@ -62,8 +70,10 @@ ) from ..projection_gathers.gathers import MuseSparkOpenTepGatedResidual, _require_aligned -#: The one row count the transport is armed for: the K = 7 verify step. +#: The row count the 8-row transport is armed for: the K = 7 verify step. ROWS = 8 +#: The row counts the wide transport is armed for: two to four requests' verify steps. +WIDE_ROWS = tuple(range(2 * ROWS, WIDE_MAX_ROWS + 1, ROWS)) def _publish_copy(slot: torch.Tensor, product: torch.Tensor) -> None: @@ -105,7 +115,16 @@ def _add_after_rmsnorm_and_norm( branch.dim() == 2 and branch.shape[1] == _unpool.local_columns(workspace.world) and residual.shape[-1] == _unpool.HIDDEN - and self._fits_fused(residual, branch, next_norm, gain, alpha, beta) + # Above 8 rows the fused join takes the wide inbox's row capacity. + and self._fits_fused( + residual, + branch, + next_norm, + gain, + alpha, + beta, + self._gather_inbox_wide if rows > ROWS else None, + ) ): raise RuntimeError( "[musespark_open] unpool_direct: the post-FFN join was armed for the direct transport but the call " @@ -152,10 +171,11 @@ def _add_after_rmsnorm_and_norm( class DirectUnpool: - """What ``DecodeTarget`` holds under ``unpool_into_boundary``: the workspace and the sites it arms. + """What ``DecodeTarget`` holds under ``unpool_into_boundary``: the workspaces and the sites it arms. ``sites`` is ``(dispatcher, join, buffer)`` per fused layer, in call order: the layer's unpool dispatcher, its post-FFN join and the transport buffer it uses. + ``wide_workspace`` (``wide_prologue``) serves the same sites at ``WIDE_ROWS``. Raises: ValueError: if two neighbouring calls of the cycle (the last and the first @@ -166,26 +186,31 @@ def __init__( self, workspace: _direct.UnpoolDirectWorkspace, sites: Sequence[Tuple[DenseGemmDispatcher, MuseSparkOpenTepDirectGatedResidual, int]], + wide_workspace: Optional[_direct.UnpoolDirectWorkspace] = None, ) -> None: self.workspace = workspace + self.wide_workspace = wide_workspace self.sites = tuple(sites) _direct.check_buffer_sequence([buffer for _, _, buffer in self.sites]) - @staticmethod - def takes(rows: int) -> bool: - """Whether a decode forward of ``rows`` rows runs on the direct transport. + def takes(self, rows: int) -> bool: + """Whether a decode forward of ``rows`` rows runs on a direct transport. - Exactly 8 rows, inside a CUDA graph capture: every rank captures the same - graphs at the same row counts, so every rank arms or none does. + Inside a CUDA graph capture: exactly 8 rows, or ``WIDE_ROWS`` with the wide + workspace. Every rank captures the same graphs at the same row counts and + installs the same workspaces, so every rank arms or none does. """ - return rows == ROWS and torch.cuda.is_current_stream_capturing() + if not torch.cuda.is_current_stream_capturing(): + return False + return rows == ROWS or (self.wide_workspace is not None and rows in WIDE_ROWS) @contextlib.contextmanager def armed(self, rows: int): """Arm every fused layer for one forward of ``rows`` rows; disarm on any exit.""" + workspace = self.workspace if rows <= ROWS else self.wide_workspace for dispatcher, join, buffer in self.sites: - dispatcher.publish = self.workspace.slot(buffer, rows) - join._unpool_direct_armed = (self.workspace, buffer) + dispatcher.publish = workspace.slot(buffer, rows) + join._unpool_direct_armed = (workspace, buffer) try: yield finally: @@ -278,10 +303,11 @@ def install( if reason is not None: raise ValueError(reason) buffers = _direct.buffer_sequence(len(fused)) - variants = sorted( - { + + def variants_at(rows: int) -> set: + return { _direct.variant( - ROWS, + rows, True, layer.post_ffn_norm.variance_epsilon, layers[index + 1].pre_qkv_proj_norm.variance_epsilon, @@ -292,13 +318,25 @@ def install( for index, layer, join, _ in fused for capture in ((False, True) if index in captured else (False,)) } - ) + + variants = sorted(variants_at(ROWS)) workspace = _direct.shared_workspace(layers[0].mapping) + # The wide layout when wide_prologue gave every fused join its wide unpool inbox and + # the layers the wide_direct_transports lever (TP4; construction-time state agreed + # over the TP group by wide_prologue.install, the same on every rank). + wide = None + if all( + join._gather_inbox_wide is not None and getattr(layer.feed_forward, "_wide_direct", False) + for _, layer, join, _ in fused + ): + wide = _direct.shared_workspace(layers[0].mapping, max_rows=WIDE_MAX_ROWS) # The compile is rank-local; its outcome is agreed too, since the collectives of # the installs after this one follow. failed: Optional[Exception] = None try: workspace.build_variants(variants) + if wide is not None: + wide.build_variants(sorted(set().union(*(variants_at(rows) for rows in WIDE_ROWS)))) except Exception as error: # noqa: BLE001 - every failure must reach the agreement failed = error reason = _tp_agreement.agreed_refusal( @@ -315,7 +353,14 @@ def install( join.__class__ = MuseSparkOpenTepDirectGatedResidual join._unpool_direct_armed = None sites.append((dispatcher, join, buffer)) - direct = DirectUnpool(workspace, sites) + direct = DirectUnpool(workspace, sites, wide) workspace.start_guard_watch() - _direct.log_once_enabled(len(sites), describe(plan.kernels).split(", ")[-1], buffers) + if wide is not None: + wide.start_guard_watch() + _direct.log_once_enabled( + len(sites), + describe(plan.kernels).split(", ")[-1], + buffers, + WIDE_ROWS if wide is not None else (), + ) return direct diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py index a29634daca08..1ef9d06da22d 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/wide_prologue/prologue.py @@ -32,6 +32,19 @@ differ from the replicated GEMM's in FP32 summation order (each output one BF16 rounding of its full-K FP32 dot product either way). +Two levers, given to ``install`` by the plan (the same on every rank): + +* ``wide_router_per_row``: the router shard of these steps is + ``router_shard_rows``, which gives each row its one-request logits (those the built-in + router gives the row when its request decodes alone) in one launch without the split-K + reduce. Those are other bits than the built-in 16..32-row launch's, so the routing of + a wide step may differ from the built-in path's. +* ``wide_direct_transports``: ``pool_into_route`` and ``unpool_into_boundary`` also build their + wide direct transports (the pool and unpool GEMMs and the router shard publish into + the routed kernel's and the boundary's transports on 16..32-row CUDA-graph steps, and on + a half-window layer the routed kernel places). Same bits as the tagged wide path given + the same router logits. + Every rank takes the same path at every call: the choice reads shapes, dtypes, contiguity and construction-time state only. """ @@ -79,12 +92,13 @@ def refusal(model: MuseSparkOpenModel): return None -def install(model: MuseSparkOpenModel) -> int: +def install(model: MuseSparkOpenModel, router_rows: bool = False, direct: bool = False) -> int: """Build the wide route and unpool inboxes and hand them to every layer. Called after ``moe_prologue.install`` and ``install_router_fork``. The refusal is agreed over the TP group before the inboxes (fabric memory, collective) are built. - Returns the number of layers installed. + ``router_rows`` (``wide_router_per_row``) and ``direct`` (``wide_direct_transports``) + are the plan's levers. Returns the number of layers installed. Raises: ValueError: on every rank alike, naming the refusal. @@ -100,15 +114,24 @@ def install(model: MuseSparkOpenModel) -> int: unpool_inbox = ProjectionInbox( mapping, row_words=_unpool.local_columns(WORLD) // 2, max_rows=WIDE_MAX_ROWS ) + router_rows, direct = bool(router_rows), bool(direct) for layer in layers: moe = layer.feed_forward moe._wide_inbox = route_inbox moe._defer_max_rows = WIDE_MAX_ROWS + moe._wide_router_rows = router_rows + moe._wide_direct = direct layer.add_ffn._gather_inbox_wide = unpool_inbox + router = ( + "router_shard_rows, each row's one-request bits" + if router_rows + else "router_shard_wide, the built-in bits" + ) logger.info( f"[musespark_open] wide_prologue: TP{WORLD} sharded MoE prologue (router shard of " - f"{_shard.local_experts(WORLD)} experts, routing in the pool kernel) and fused unpool " - f"gather on decode steps of {_shard.WIDE_MIN_ROWS}..{_shard.WIDE_MAX_ROWS} rows " - f"on all {len(layers)} layers" + f"{_shard.local_experts(WORLD)} experts ({router}), routing in the pool kernel) and fused " + f"unpool gather on decode steps of {_shard.WIDE_MIN_ROWS}..{_shard.WIDE_MAX_ROWS} rows " + f"on all {len(layers)} layers; wide direct transports " + f"{'on' if direct else 'off'} (wide_direct_transports)" ) return len(layers) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/post_attn_hop.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/post_attn_hop.py index 43c19e7db670..2d906f96bf8e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/post_attn_hop.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/allreduce/post_attn_hop.py @@ -81,7 +81,9 @@ word whose two halves arrive in separate stores is not taken half-written. And a GEMM cannot read the arrival counter, so the buffer index is a launch argument the caller fixes per call site (layer ``l`` uses ``l mod 2``, see -``PostAttnDirectWorkspace``). +``PostAttnDirectWorkspace``). At TP4 a second direct transport with 32-row slots +(``max_rows=WIDE_MAX_ROWS``) takes the 9..32-row calls the same way, on the 4-CTA split +of the wide hop. What the direct transport needs from its users, and what checks it. The producer must store each output element once and never reduce through its @@ -106,7 +108,7 @@ import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -1170,7 +1172,8 @@ def hop_supported( #: #: * ``(10, 3)``: cuBLAS's choice for the 8-row o_proj, ``[8, 1024] x [1024, 8192]`` BF16. #: * ``(10, 7)``: flashinfer's split-K and direct kernels, the o_proj kernels the dense GEMM -#: plan names there; cuBLAS's choice on this arch is not one. +#: plan names there (the split-K tactics of 9..32 rows on the wide direct transport +#: included, ``wide_dense_gemms``); cuBLAS's choice on this arch is not one. VALIDATED_PRODUCERS: Dict[Tuple[int, int], FrozenSet[str]] = { (10, 3): frozenset({"nvjet_sm103_tst_64x8_64x16_4x1_v_bz_TNT"}), (10, 7): frozenset( @@ -1330,7 +1333,9 @@ class _PostAttnConsumerKernel(_PostAttnHopKernel): rank's copy. Every slot of the local copy is polled, this rank's own (the multicast loopback) included, with the stale-only schedule: one read of every slot, then reads of only the slots that still hold a sentinel half, all issued - before any is tested. Then every slot is reset. + before any is tested. Then every slot is reset. ``max_rows`` is the row capacity + of the transport's slots (its slot stride): 8, or ``WIDE_MAX_ROWS`` for the wide + direct transport, whose calls of 9..32 rows run the boundary's 4-CTA split. PDL memory order: before ``griddepcontrol.wait`` the kernel touches shared memory and the cluster barrier only. The residual, the coefficients and the @@ -1352,8 +1357,19 @@ def __init__( early_trigger: bool, guard_ns: int, guard_break: bool, + max_rows: int = MAX_ROWS, ): - super().__init__(ranks, cluster, True, post_eps, pre_eps, norm_fp32, pdl, early_trigger) + super().__init__( + ranks, + cluster, + True, + post_eps, + pre_eps, + norm_fp32, + pdl, + early_trigger, + max_rows=max_rows, + ) self.guard_ns = int(guard_ns) self.guard_break = bool(guard_break) @@ -1449,7 +1465,7 @@ def kernel( receive_address = receive.toint() values = cute.make_rmem_tensor(cute.make_layout((self.ranks, 4)), Uint32) for p in cutlass.range_constexpr(self.ranks): - address = receive_address + (buffer * MAX_RANKS + p) * SLAB_BYTES + slot + address = receive_address + (buffer * MAX_RANKS + p) * self.slab_bytes + slot v0, v1, v2, v3 = _poll_x4( address, Uint32(1), Uint32(0), Uint32(0), Uint32(0), Uint32(0) ) @@ -1468,7 +1484,9 @@ def kernel( while stale != Uint32(0): polled = stale for p in cutlass.range_constexpr(self.ranks): - address = receive_address + (buffer * MAX_RANKS + p) * SLAB_BYTES + slot + address = ( + receive_address + (buffer * MAX_RANKS + p) * self.slab_bytes + slot + ) v0, v1, v2, v3 = _poll_x4( address, (polled >> Uint32(p)) & Uint32(1), @@ -1498,7 +1516,9 @@ def kernel( stale = Uint32(0) # Retire every slot of this call before this grid can complete. for p in cutlass.range_constexpr(self.ranks): - _direct_reset_x4(receive_address + (buffer * MAX_RANKS + p) * SLAB_BYTES + slot) + _direct_reset_x4( + receive_address + (buffer * MAX_RANKS + p) * self.slab_bytes + slot + ) # Early trigger: past the last transport access (the hop's early-trigger placement). if cutlass.const_expr(self.pdl and self.early_trigger): cute.arch.griddepcontrol_launch_dependents() @@ -1594,9 +1614,19 @@ def _compiled_consumer( early_trigger: bool, guard_ns: int, guard_break: bool, + max_rows: int = MAX_ROWS, ): kernel = _PostAttnConsumerKernel( - ranks, cluster, post_eps, pre_eps, norm_fp32, pdl, early_trigger, guard_ns, guard_break + ranks, + cluster, + post_eps, + pre_eps, + norm_fp32, + pdl, + early_trigger, + guard_ns, + guard_break, + max_rows, ) with torch.cuda.device(device_index): stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) @@ -1647,6 +1677,14 @@ class PostAttnDirectWorkspace: daemon thread reads it every ``watchdog_s`` seconds without waiting for the device and logs an error once it is set (0 = no thread). ``reset`` and ``check_producer`` are the collective recovery and the producer check (module docstring). + + ``max_rows=WIDE_MAX_ROWS`` builds the wide direct transport (4 ranks only): the same + layout with slots of 32 rows, calls of ``MAX_ROWS + 1 .. WIDE_MAX_ROWS`` rows on the + boundary's 4-CTA split (``cluster_for_rows``). It is a separate allocation, so its calls + and the 8-row transport's interleave freely, and the buffer argument above holds per + transport. A producer of fewer rows than a slot holds stores only its own rows (the + flashinfer split-K epilogue predicates the M tail) and the consumer resets only those, so + calls of different row counts share the buffers. """ def __init__( @@ -1658,6 +1696,7 @@ def __init__( guard_ns: int = DIRECT_GUARD_NS, guard_break: bool = False, watchdog_s: float = DIRECT_WATCHDOG_S, + max_rows: int = MAX_ROWS, ): from tensorrt_llm._torch.distributed.ops import ( _get_mnnvl_workspace_comm, @@ -1674,6 +1713,16 @@ def __init__( self.guard_break = bool(guard_break) if self.ranks not in (4, MAX_RANKS): raise ValueError(f"the direct hop serves 4 or 8 ranks, got {self.ranks}") + if max_rows not in (MAX_ROWS, WIDE_MAX_ROWS) or ( + max_rows == WIDE_MAX_ROWS and self.ranks != 4 + ): + raise ValueError( + f"the direct hop's transport holds {MAX_ROWS} rows, or {WIDE_MAX_ROWS} at 4 " + f"ranks; got {max_rows} rows at {self.ranks} ranks" + ) + self.max_rows = int(max_rows) + #: The row counts this transport's calls take. + self.min_rows = 1 if self.max_rows == MAX_ROWS else MAX_ROWS + 1 self.device = torch.device("cuda", _mnnvl_device_index(mapping)) #: The device's compute capability, which selects ``check_producer``'s validated set. self.capability: Optional[Tuple[int, int]] = None @@ -1687,13 +1736,14 @@ def __init__( if self.capability not in ((10, 3), (10, 7)): raise RuntimeError("the direct hop runs on SM103 or SM107 only") use_fabric = mapping.is_multi_node() - self._handle = _make_mnnvl_mcast_buffer( - self._comm, WORKSPACE_BYTES, mapping, use_fabric - ) - words = WORKSPACE_BYTES // 4 + nbytes = workspace_bytes(self.max_rows) + self._handle = _make_mnnvl_mcast_buffer(self._comm, nbytes, mapping, use_fabric) + words = nbytes // 4 self._receive = self._handle.get_uc_buffer(self.rank, (words,), torch.int32, 0) self._mc = self._handle.get_mc_buffer((words,), torch.int32, 0) - self._slots = self._mc.view(torch.bfloat16).view(BUFFERS, MAX_RANKS, MAX_ROWS, HIDDEN) + self._slots = self._mc.view(torch.bfloat16).view( + BUFFERS, MAX_RANKS, self.max_rows, HIDDEN + ) with torch.cuda.device(self.device): self._receive.fill_(_DIRECT_SENTINEL_I32) # Pinned host memory: the kernel's guard store reaches the host directly. @@ -1718,16 +1768,22 @@ def __init__( ).start() def build_variants(self, variants: Sequence[Tuple[float, float, bool]]) -> None: - """Compile and load both splits of each ``(post_eps, pre_eps, norm_fp32)`` not built yet. + """Compile and load every split of each ``(post_eps, pre_eps, norm_fp32)`` not built yet. - Collective: either every rank of the group has them afterwards or every rank raises. + The splits some row count of this transport launches: 16 and 8 on the 8-row one, 4 on + the wide one. Collective: either every rank of the group has them afterwards or every + rank raises. """ from tensorrt_llm._torch.distributed.ops import _mnnvl_workspace_all_succeeded error: Optional[BaseException] = None try: + clusters = sorted( + {cluster_for_rows(rows) for rows in range(self.min_rows, self.max_rows + 1)}, + reverse=True, + ) for post_eps, pre_eps, norm_fp32 in variants: - for cluster in (16, 8): + for cluster in clusters: self._kernel(cluster, post_eps, pre_eps, norm_fp32) torch.cuda.synchronize(self.device) except Exception as exc: # noqa: BLE001 - every failure must reach the agreement @@ -1758,6 +1814,7 @@ def _kernel(self, cluster: int, post_eps: float, pre_eps: float, norm_fp32: bool self.early_trigger, self.guard_ns, self.guard_break, + self.max_rows, ) # A dry launch loads the module now, not inside a capture. dummy = self.control.data_ptr() @@ -1772,8 +1829,9 @@ def slot(self, parity: int, rows: int) -> torch.Tensor: """This rank's ``[rows, 8192]`` BF16 slot of buffer ``parity``, on the multicast address. Write-only: a store through it reaches every rank's copy; nothing may read it. + ``rows`` is a row count this transport's calls take (``min_rows .. max_rows``). """ - if parity not in (0, 1) or not 1 <= rows <= MAX_ROWS: + if parity not in (0, 1) or not self.min_rows <= rows <= self.max_rows: raise ValueError(f"no slot for parity {parity}, {rows} rows") return self._slots[parity, self.rank, :rows] @@ -1904,16 +1962,22 @@ def shared_direct_workspace( *, variants: Sequence[Tuple[float, float, bool]] = (), early_trigger: bool = False, + max_rows: int = MAX_ROWS, ) -> PostAttnDirectWorkspace: - """The one direct workspace of this group and trigger placement, built on first use. + """The one direct workspace of this group, trigger placement and row capacity, built on + first use. Collective, like the constructor; a later call gets the same workspace back, with any - variant it names built as well. + variant it names built as well. ``max_rows=WIDE_MAX_ROWS`` asks for the wide one. """ key = (mapping, bool(early_trigger), _LAUNCH_PDL) + if max_rows != MAX_ROWS: + key = key + (int(max_rows),) workspace = _SHARED_DIRECT.get(key) if workspace is None: - workspace = PostAttnDirectWorkspace(mapping, variants=variants, early_trigger=early_trigger) + workspace = PostAttnDirectWorkspace( + mapping, variants=variants, early_trigger=early_trigger, max_rows=max_rows + ) _SHARED_DIRECT[key] = workspace else: workspace.build_variants(variants) @@ -1929,10 +1993,12 @@ def consumer_hop_supported( parity: int, workspace: PostAttnDirectWorkspace, ) -> bool: - """``hop_supported`` for the consumer: ``published`` must be ``workspace.slot(parity, rows)``.""" + """``hop_supported`` for the consumer: ``published`` must be ``workspace.slot(parity, rows)`` + for a row count the workspace's calls take (``min_rows .. max_rows``).""" rows = published.shape[0] if published.ndim == 2 else 0 return ( - hop_takes(published, residual, alpha, beta, next_weight) + workspace.min_rows <= rows + and hop_takes(published, residual, alpha, beta, next_weight, workspace.max_rows) and hop_layout_ok(residual, alpha, beta, next_weight) and parity in (0, 1) and published.data_ptr() == workspace.slot(parity, rows).data_ptr() @@ -1950,17 +2016,66 @@ class PostAttnConsumerHop(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell(why="no speculation: one row, the 16-CTA cluster split", spec=dict(rows=1)), + Cell(why="two rows: the smallest 8-CTA split", spec=dict(rows=2)), + Cell(why="K = 3 verify width", spec=dict(rows=4)), + Cell(why="K = 5 verify width", spec=dict(rows=6)), + Cell(why="K = 7 verify width, the one row count the target serves", spec=dict(rows=8)), + Cell( + why="the second buffer: odd layers publish into parity 1", + spec=dict(rows=8, parity=1), + ), + Cell( + why="the FP32-residual prenorm ordering (MUSESPARK_OPEN_PREQKV_NORM_FP32_RESIDUAL)", + spec=dict(rows=8, norm_fp32=True), + ), + Cell( + why="exact zeros and negative zeros published unchanged (0x80000000 is a payload here)", + spec=dict(rows=8, zeros=True), + ), + Cell( + why="cancelling partials: the FP32 rank-order sum differs from other orders", + spec=dict(rows=8, cancel=True), + ), + Cell( + why="wide direct transport, two requests at K = 7: the 4-CTA split", + spec=dict(rows=16, wide=True), + ), + Cell( + why="wide direct transport, three requests at K = 7, the second buffer", + spec=dict(rows=24, wide=True, parity=1), + ), + Cell(why="wide direct transport, four requests at K = 7", spec=dict(rows=32, wide=True)), + Cell( + why="wide direct transport, the FP32-residual prenorm ordering", + spec=dict(rows=32, wide=True, norm_fp32=True), + ), + Cell( + why="wide direct transport, exact zeros and negative zeros", + spec=dict(rows=24, wide=True, zeros=True), + ), + Cell( + why="wide direct transport, cancelling partials", + spec=dict(rows=32, wide=True, cancel=True), + ), + ) + note = """ Metadata consumed: none. Call site: once per decoder layer per decode step of - the row count the target serves directly (8 at K = 7), from the target's - attention-side residual join, after the layer's o_proj GEMM wrote its output - into ``published``; layer ``l`` uses parity ``l mod 2``. + the row counts the target serves directly (8 at K = 7; with + ``wide_oproj_publish`` also 16, 24 and 32 on the wide direct + transport at TP4), from the target's attention-side residual join, after the layer's + o_proj GEMM wrote its output into ``published``; layer ``l`` uses parity ``l mod 2``. - Bit-identical to ``post_attn_hop`` on the same partials, gated by its own + Bit-identical to ``post_attn_hop`` on the same partials (on the wide direct + transport: to ``post_attn_hop`` on its wide transport), gated by its own multi-GPU test (both outputs, every cell, a real cuBLAS o_proj as the - producer, graph replay with changing inputs, a delayed rank, a long burst, - a concurrent poller that records the first non-sentinel value of every - word). Ranks 4 and 8. The guard word turns a broken call sequence into a + producer and, on the wide transport, flashinfer's wide split-K o_proj tactics, + graph replay with changing inputs, a delayed rank, a long burst, calls of mixed + row counts on one wide buffer, a concurrent poller that records the first + non-sentinel value of every word). Ranks 4 and 8 (the wide transport at 4). The + guard word turns a broken call sequence into a flag instead of only a hang (tested with ``guard_break``), and the workspace's watchdog logs it while the poll still waits. A sequence broken between a producer and this op is cleared by the workspace's collective @@ -1968,7 +2083,8 @@ class PostAttnConsumerHop(OpWrapper): Loud preconditions: `consumer_hop_supported` must hold. Not certified: producers that reduce through their output (split-K through global memory), - rows above 8, widths other than 8192; ``check_producer`` accepts only the device's + rows above 8 on the 8-row transport or above 32 on the wide one, the wide transport + at 8 ranks, widths other than 8192; ``check_producer`` accepts only the device's ``VALIDATED_PRODUCERS``. """ @@ -1989,7 +2105,8 @@ def __call__( published, residual, alpha, beta, next_weight, parity, workspace ): raise ValueError( - "post_attn_consumer_hop takes this rank's [1..8, 8192] slot of buffer " + f"post_attn_consumer_hop takes this rank's [{workspace.min_rows}.." + f"{workspace.max_rows}, 8192] slot of buffer " "`parity` of its workspace as `published`, an FP32 residual of the same " "shape, FP32 alpha/beta and a BF16 weight of 8192, 16-byte aligned; got " f"{published.dtype} {tuple(published.shape)} / {residual.dtype} " @@ -2049,6 +2166,18 @@ def log_once_direct( ) +def log_once_wide_direct( + layers: int, rows: Sequence[int], kernels: str, prefetch_dropped: int = 0 +) -> None: + logger.info( + "[musespark_open] attention_output_into_hop: o_proj also publishes into the " + f"post-attention hop's {WIDE_MAX_ROWS}-row direct transport on the decode steps of " + f"{'/'.join(str(r) for r in rows)} rows, all {layers} layers (o_proj {kernels}; the " + "wide hop's own publication gone; the o_proj -> pool-weight prefetch dropped at those " + f"rows on {prefetch_dropped} layers)" + ) + + def log_once_producer(name: str) -> None: logger.info( f"[musespark_open] attention_output_into_hop: the published o_proj captures {name}, a validated producer" @@ -2067,6 +2196,39 @@ class PostAttnHop(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + # The decode step of this target at K = 7 carries 8 rows; K = 3 / 5 carry 4 / 6; + # one row is the no-speculation step and the only shape on the 16-CTA split. + CELLS: tuple[Cell, ...] = ( + Cell(why="no speculation: one row, the 16-CTA cluster split", spec=dict(rows=1)), + Cell(why="two rows: the smallest 8-CTA split", spec=dict(rows=2)), + Cell(why="K = 3 verify width", spec=dict(rows=4)), + Cell(why="K = 5 verify width", spec=dict(rows=6)), + Cell(why="K = 7 verify width, the row count the target serves", spec=dict(rows=8)), + Cell( + why="the FP32-residual prenorm ordering (MUSESPARK_OPEN_PREQKV_NORM_FP32_RESIDUAL)", + spec=dict(rows=8, norm_fp32=True), + ), + Cell( + why="exact zeros and negative zeros in every rank's partial", + spec=dict(rows=8, zeros=True), + ), + Cell( + why="cancelling partials: the FP32 rank-order sum differs from other orders", + spec=dict(rows=8, cancel=True), + ), + Cell( + why="wide transport, two requests at K = 7: the 4-CTA split", + spec=dict(rows=16, wide=True), + ), + Cell(why="wide transport, the smallest wide call", spec=dict(rows=9, wide=True)), + Cell(why="wide transport, three requests at K = 7", spec=dict(rows=24, wide=True)), + Cell(why="wide transport, four requests at K = 7", spec=dict(rows=32, wide=True)), + Cell( + why="wide transport, cancelling partials", + spec=dict(rows=32, wide=True, cancel=True), + ), + ) + note = """ Metadata consumed: none. Call site: once per decoder layer per decode step of at most 8 rows, from the target's attention-side residual join, in place diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_direct_route.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_direct_route.py index 8b5582052f47..1b18141b307b 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_direct_route.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/pool_direct_route.py @@ -61,6 +61,14 @@ records, after its own dependency wait. A poll that has seen no arrival for ``guard_ns`` sets a guard word in host memory (and keeps polling). +Wide transport (TP4, ``max_rows=WIDE_MAX_ROWS``): the decode steps of two to four +requests (16..32 rows, ``wide_prologue``) run the wide tagged kernel's statements, in +which each CTA routes only its own two rows (no placement, no early ids). A record is +then the source rank's ``[32, C]`` pool slot followed by one ``[32, L]`` FP32 logit +block, the router's output as it is (``router_shard_rows`` with the block as its +output): CTA ``c`` polls and resets rows ``2c`` and ``2c + 1`` of each sender's block. +A separate workspace with its own buffer sequence; the 8-row layout is unchanged. + PDL memory order: before ``griddepcontrol.wait`` the routing warps read the expert bias (a parameter), as in ``pool_gather_route``; nothing else, in particular nothing of the transport. After it: the transport (ordered by the @@ -84,7 +92,7 @@ import torch from cuda.bindings import runtime as _cudart -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -92,7 +100,7 @@ from . import pool_gather_route as _route from . import unpool_direct_boundary as _direct from .pool_gather_fp4quant import _VECTOR_SPAN_WORDS, _VECTORS, BLOCK, HIDDEN, local_columns -from .projection_inbox import MAX_ROWS +from .projection_inbox import MAX_ROWS, WIDE_MAX_ROWS #: The TP sizes this transport is laid out for: TP8 over two trays and TP4 on one. WORLDS = (4, 8) @@ -118,6 +126,11 @@ #: CTAs of the routed kernel at ``MAX_ROWS`` rows (two rows each): one logit box each. CTAS = MAX_ROWS // 2 CHUNKS = _route.CHUNKS +#: The wide layout's half-window census (``PLACE_HALF``), per buffer in rank-local device +#: memory: ``CHUNKS`` presence words, the CTAs' arrival count and their read count. +CENSUS_WORDS = 2 * CHUNKS +CENSUS_ARRIVED = CHUNKS +CENSUS_READ = CHUNKS + 1 def pool_row_words(world: int) -> int: @@ -125,24 +138,39 @@ def pool_row_words(world: int) -> int: return local_columns(world) // 2 -def pool_slot_words(world: int) -> int: - """32-bit words of a pool slot: ``MAX_ROWS`` rows.""" - return MAX_ROWS * pool_row_words(world) +def check_layout(world: int, max_rows: int) -> None: + """Raise unless ``(world, max_rows)`` is a served layout: 8 rows at TP8 or TP4, 32 rows at TP4.""" + if world not in WORLDS or max_rows not in (MAX_ROWS, WIDE_MAX_ROWS): + raise ValueError(f"no direct pool layout for TP{world} and {max_rows} rows") + if max_rows == WIDE_MAX_ROWS and world != 4: + raise ValueError("the wide direct pool layout is TP4 only") + + +def pool_slot_words(world: int, max_rows: int = MAX_ROWS) -> int: + """32-bit words of a pool slot: ``max_rows`` rows.""" + return max_rows * pool_row_words(world) def box_words(world: int) -> int: - """32-bit words of a logit box: ``[MAX_ROWS, local_experts(world)]`` FP32.""" + """32-bit words of a logit box of the 8-row layout: ``[MAX_ROWS, local_experts(world)]`` FP32.""" return MAX_ROWS * _route.local_experts(world) -def record_words(world: int) -> int: - """32-bit words of a record: the pool slot, then ``CTAS`` logit boxes.""" - return pool_slot_words(world) + CTAS * box_words(world) +def logit_words(world: int, max_rows: int = MAX_ROWS) -> int: + """32-bit words of a record's logits: ``CTAS`` boxes (8 rows), or one ``[max_rows, L]`` block (wide).""" + if max_rows == MAX_ROWS: + return CTAS * box_words(world) + return max_rows * _route.local_experts(world) + + +def record_words(world: int, max_rows: int = MAX_ROWS) -> int: + """32-bit words of a record: the pool slot, then the logits (``logit_words``).""" + return pool_slot_words(world, max_rows) + logit_words(world, max_rows) -def workspace_bytes(world: int) -> int: +def workspace_bytes(world: int, max_rows: int = MAX_ROWS) -> int: """Bytes of one rank's copy of the transport: ``BUFFERS x world`` records.""" - return BUFFERS * world * record_words(world) * 4 + return BUFFERS * world * record_words(world, max_rows) * 4 #: 32-bit words of one row of a pool slot, of a pool slot, of a logit box, of a record @@ -213,9 +241,71 @@ def workspace_bytes(world: int) -> int: _stamp_timer, _top8_by_redux, ) - from .projection_inbox import _asm, _u32 + from .projection_inbox import _asm, _i64, _u32 from .unpool_direct_boundary import _globaltimer, _pending, _poll_x4, _reset_x4, _set_guard_word + @dsl_user_op + def _red_global_or(address, value, *, loc=None, ip=None): + """A relaxed GPU-scope OR into a rank-local global word (the wide census).""" + _asm( + None, + [_i64(address, loc, ip), _u32(value, loc, ip)], + "red.relaxed.gpu.global.or.b32 [$0], $1;", + "l,r", + loc, + ip, + ) + + @dsl_user_op + def _red_release_gpu_add(address, value, *, loc=None, ip=None): + """A release GPU-scope add (a CTA's arrival at the wide census).""" + _asm( + None, + [_i64(address, loc, ip), _u32(value, loc, ip)], + "red.release.gpu.global.add.u32 [$0], $1;", + "l,r", + loc, + ip, + ) + + @dsl_user_op + def _atom_acq_rel_gpu_add(address, value, *, loc=None, ip=None): + """An acquire-release GPU-scope add returning the old value.""" + return Uint32( + _asm( + T.i32(), + [_i64(address, loc, ip), _u32(value, loc, ip)], + "atom.acq_rel.gpu.global.add.u32 $0, [$1], $2;", + "=r,l,r", + loc, + ip, + ) + ) + + @dsl_user_op + def _ld_acquire_gpu_b32(address, *, loc=None, ip=None): + return Uint32( + _asm( + T.i32(), + [_i64(address, loc, ip)], + "ld.acquire.gpu.global.b32 $0, [$1];", + "=r,l", + loc, + ip, + ) + ) + + @dsl_user_op + def _st_relaxed_gpu_b32(address, value, *, loc=None, ip=None): + _asm( + None, + [_i64(address, loc, ip), _u32(value, loc, ip)], + "st.relaxed.gpu.global.b32 [$0], $1;", + "l,r", + loc, + ip, + ) + @dsl_user_op def _pending_words(w0, w1, w2, w3, *, loc=None, ip=None): """1 while any of the four 32-bit words is ``DIRECT_SENTINEL``, else 0.""" @@ -256,10 +346,24 @@ def __init__( guard_ns: int, guard_break: bool, world: int = WORLD, + max_rows: int = MAX_ROWS, + census: int = 0, ): + check_layout(world, max_rows) + wide = max_rows > MAX_ROWS + if wide and ( + place not in (PLACE_NONE, PLACE_HALF) or (place == PLACE_HALF) != bool(census) + ): + raise ValueError( + "the wide layout routes only, or places the half window with its census words" + ) + # A wide layout compiles the tagged kernel's wide statements (``self.wide``): + # each CTA routes its own rows, without early ids; the tagged kernel has no + # wide placement, so it is built routing-only and the half window's placement, + # over the grid's census (``census``), is this kernel's. super().__init__( peer_rank, - place, + PLACE_NONE if wide else place, static_from, True, False, @@ -268,13 +372,16 @@ def __init__( trigger, stamps, world, + max_rows, ) + self.place = place + self.census = int(census) self.guard_ns = int(guard_ns) self.guard_break = bool(guard_break) self.pool_row_words = pool_row_words(world) - self.pool_slot_words = pool_slot_words(world) + self.pool_slot_words = pool_slot_words(world, max_rows) self.box_words = box_words(world) - self.record_words = record_words(world) + self.record_words = record_words(world, max_rows) @cute.jit def __call__( @@ -586,11 +693,22 @@ def kernel( if cutlass.const_expr(self.place != PLACE_NONE): if rt < CHUNKS: _st_shared_b32(sCensus.iterator + rt, Uint32(0)) - if rt * 4 < M * local_count: + if cutlass.const_expr(self.wide): + # The CTA's own rows of each sender's block: words 4 rt of its + # [rows_per_block, L] rows (fewer when its last row is past M). + word_limit = _min_u32( + Uint32(rows_per_block * local_count), + Uint32((M - bidx * rows_per_block) * local_count), + ) + logits_start = Int64(self.pool_slot_words) + Int64( + bidx * rows_per_block * local_count + ) + else: + word_limit = M * local_count + logits_start = Int64(self.pool_slot_words) + Int64(bidx * self.box_words) + if rt * 4 < word_limit: word = Int64(rt * 4) - piece = ( - Int64(self.pool_slot_words) + Int64(bidx * self.box_words) + word - ) * 4 + piece = (logits_start + word) * 4 received = cute.make_rmem_tensor(cute.make_layout((4 * world,)), Uint32) start = Int64(0) if cutlass.const_expr(self.guard_ns > 0): @@ -641,13 +759,20 @@ def kernel( full_row = bidx * rows_per_block + rt // 64 if full_row < M: lr = full_row * EXPERTS + (rt % 64) * 4 - f0, f1, f2, f3 = _load_shared_words(sLogits.iterator + lr) + if cutlass.const_expr(self.wide): + slr = (rt // 64) * EXPERTS + (rt % 64) * 4 + else: + slr = lr + f0, f1, f2, f3 = _load_shared_words(sLogits.iterator + slr) _store_global_words(logits_full + Int64(lr) * 4, f0, f1, f2, f3) eps_bits = _f32_bits(renorm_eps) # Each lane's keys of the warp's two rows, sorted; then the warp's - # top 8 of each row (key s in lane s). + # top 8 of each row (key s in lane s). Wide: one row per routing warp, + # the CTA's own (warps 0 and 1; warps 2 and 3 route unread shared rows + # and store nothing). + halves = 1 if self.wide else 2 lists = [] - for half in cutlass.range_constexpr(2): + for half in cutlass.range_constexpr(halves): row = rwarp + 4 * half keys = [] for k in cutlass.range_constexpr(CHUNKS): @@ -658,12 +783,22 @@ def kernel( ) lists.append(_sort8(keys)) picks, winners = _top8_by_redux(lists, lane) - for half in cutlass.range_constexpr(2): + for half in cutlass.range_constexpr(halves): row = rwarp + 4 * half if cutlass.const_expr(self.place != PLACE_NONE): # The census, from the rounds' winners (warp-uniform): one - # lane sets each selected expert's bit. - counts = _eq(Uint32(lane), Uint32(0)) & _lt_u(Uint32(row), Uint32(M)) + # lane sets each selected expert's bit (wide: of the CTA's own + # rows only; the grid's census is merged below). + if cutlass.const_expr(self.wide): + counts = ( + _eq(Uint32(lane), Uint32(0)) + & _lt_u(Uint32(row), Uint32(rows_per_block)) + & _lt_u(Uint32(bidx * rows_per_block + row), Uint32(M)) + ) + else: + counts = _eq(Uint32(lane), Uint32(0)) & _lt_u( + Uint32(row), Uint32(M) + ) if counts != Uint32(0): for s in cutlass.range_constexpr(TOP_K): chosen = Uint32(255) - winners[half][s] @@ -671,7 +806,7 @@ def kernel( sCensus.iterator + (chosen >> 5), Uint32(1) << (chosen & Uint32(31)), ) - for half in cutlass.range_constexpr(2): + for half in cutlass.range_constexpr(halves): row = rwarp + 4 * half mine = picks[half] # Masked so that the unused rows' (unwritten) keys stay in bounds. @@ -683,19 +818,81 @@ def kernel( quad = _add_rn(pair, _shfl_xor_b32(pair, 4)) total = _add_rn(quad, _shfl_xor_b32(quad, 2)) weight = _div_rn(score, _add_rn(eps_bits, total)) - keep = _lt_u(Uint32(row), Uint32(M)) & _lt_u(Uint32(lane), Uint32(TOP_K)) - if keep != Uint32(0): - _st_shared_b32(sSelected.iterator + row * TOP_K + lane, expert_id) - _st_shared_b32(sWeights.iterator + row * TOP_K + lane, weight) - if row // rows_per_block == bidx: - out = Int64(row * TOP_K + lane) * 4 + if cutlass.const_expr(self.wide): + global_row = bidx * rows_per_block + row + keep = ( + _lt_u(Uint32(row), Uint32(rows_per_block)) + & _lt_u(Uint32(global_row), Uint32(M)) + & _lt_u(Uint32(lane), Uint32(TOP_K)) + ) + if keep != Uint32(0): + if cutlass.const_expr(self.place != PLACE_NONE): + # The CTA's own rows, at their CTA-relative rows. + _st_shared_b32( + sSelected.iterator + row * TOP_K + lane, expert_id + ) + _st_shared_b32(sWeights.iterator + row * TOP_K + lane, weight) + out = Int64(global_row * TOP_K + lane) * 4 _st_global_b32(selected + out, expert_id) _st_global_b32(weights + out, weight) + else: + keep = _lt_u(Uint32(row), Uint32(M)) & _lt_u( + Uint32(lane), Uint32(TOP_K) + ) + if keep != Uint32(0): + _st_shared_b32(sSelected.iterator + row * TOP_K + lane, expert_id) + _st_shared_b32(sWeights.iterator + row * TOP_K + lane, weight) + if row // rows_per_block == bidx: + out = Int64(row * TOP_K + lane) * 4 + _st_global_b32(selected + out, expert_id) + _st_global_b32(weights + out, weight) # Placement (replicated experts) from all M rows' selections, in every # routing warp. if cutlass.const_expr(self.place != PLACE_NONE): _bar_sync(ROUTE_BARRIER, ROUTE_THREADS) + if cutlass.const_expr(self.wide): + # The grid's census: every CTA ORs its rows' bits into this + # rank's census words of the buffer, arrives, and waits for + # every CTA's arrival (all of the grid's CTAs are resident: + # at most 16); then each CTA reads the merged words. The + # last CTA to have read them clears them for the next call. + census = Int64(self.census) + Int64(buffer) * Int64(CENSUS_WORDS * 4) + if rt < CHUNKS: + bits = _ld_shared_b32(sCensus.iterator + rt) + if bits != Uint32(0): + _red_global_or(census + Int64(rt * 4), bits) + _bar_sync(ROUTE_BARRIER, ROUTE_THREADS) + ctas = Uint32((M + rows_per_block - 1) // rows_per_block) + if rt == 0: + _red_release_gpu_add(census + Int64(CENSUS_ARRIVED * 4), Uint32(1)) + start = Int64(0) + if cutlass.const_expr(self.guard_ns > 0): + start = _globaltimer() + arrived = _ld_acquire_gpu_b32(census + Int64(CENSUS_ARRIVED * 4)) + while _lt_u(arrived, ctas) != Uint32(0): + arrived = _ld_acquire_gpu_b32( + census + Int64(CENSUS_ARRIVED * 4) + ) + if cutlass.const_expr(self.guard_ns > 0): + if _globaltimer() - start > Int64(self.guard_ns): + _set_guard_word(control + Int64(GUARD_WORD * 4)) + if cutlass.const_expr(self.guard_break): + arrived = ctas + _bar_sync(ROUTE_BARRIER, ROUTE_THREADS) + if rt < CHUNKS: + _st_shared_b32( + sCensus.iterator + rt, + _ld_relaxed_gpu_b32(census + Int64(rt * 4)), + ) + _bar_sync(ROUTE_BARRIER, ROUTE_THREADS) + if rt == 0: + read = _atom_acq_rel_gpu_add( + census + Int64(CENSUS_READ * 4), Uint32(1) + ) + if _eq(read + Uint32(1), ctas) != Uint32(0): + for word in cutlass.range_constexpr(CENSUS_WORDS): + _st_relaxed_gpu_b32(census + Int64(word * 4), Uint32(0)) w0, w1, w2, w3 = _load_shared_words(sCensus.iterator) w4, w5, w6, w7 = _load_shared_words(sCensus.iterator + 4) words = [w0, w1, w2, w3, w4, w5, w6, w7] @@ -848,12 +1045,14 @@ def kernel( ) if writes != Uint32(0): slot_index = own_row * TOP_K + lane - g = _ld_shared_b32(sSelected.iterator + slot_index) + # Wide: the CTA's rows sit at their CTA-relative rows. + shared_index = rwarp * TOP_K + lane if self.wide else slot_index + g = _ld_shared_b32(sSelected.iterator + shared_index) placed = _placed_expert( g, words, xs, self.peer_rank, x2s, world, place_half ) _st_global_b32(local_ids + Int64(slot_index) * 4, placed) - w_bits = _ld_shared_b32(sWeights.iterator + slot_index) + w_bits = _ld_shared_b32(sWeights.iterator + shared_index) _st_global_b16(weights_bf16 + Int64(slot_index) * 2, _bf16_bits(w_bits)) held_load = loads[0] for b in cutlass.range_constexpr(1, world): @@ -889,6 +1088,8 @@ def _compile( guard_ns: int, guard_break: bool, world: int = WORLD, + max_rows: int = MAX_ROWS, + census: int = 0, ): kernel = _PoolDirectRouteKernel( peer_rank, @@ -901,6 +1102,8 @@ def _compile( guard_ns, guard_break, world, + max_rows, + census, ) rows = cute.sym_int() fake = cute.runtime.make_fake_compact_tensor @@ -980,9 +1183,10 @@ class PoolDirectWorkspace: ``slot(buffer, rows)`` is this rank's ``[rows, local_columns(world)]`` BF16 pool slot of ``buffer`` (512 columns at TP8, 1024 at TP4) and ``boxes(buffer, rows)`` - its ``[publish_copies(rows), box_words(world)]`` FP32 logit boxes, both on the - multicast address: a producer given them as its output - publishes into every rank's copy. ``kernel(variant)`` is the compiled routed + its ``[publish_copies(rows), box_words(world)]`` FP32 logit boxes (the 8-row + layout) or ``block(buffer, rows)`` its ``[rows, L]`` FP32 logit block (the wide + layout, ``max_rows=WIDE_MAX_ROWS``, TP4), all on the multicast address: a producer + given them as its output publishes into every rank's copy. ``kernel(variant)`` is the compiled routed kernel for a call's settings. Construction is collective over the group (fabric allocation and sentinel fill), and so is ``build_variants``; both succeed on every rank or raise on every rank. Nothing compiles inside a CUDA graph @@ -999,6 +1203,7 @@ def __init__( *, guard_ns: int = DIRECT_GUARD_NS, guard_break: bool = False, + max_rows: int = MAX_ROWS, ): from tensorrt_llm._torch.distributed.ops import ( _get_mnnvl_workspace_comm, @@ -1011,10 +1216,13 @@ def __init__( raise ValueError( "the direct pool transport needs TP8 or TP4 without attention DP or CP" ) + check_layout(int(mapping.tp_size), int(max_rows)) self.rank = mapping.tp_rank - #: The group size, which fixes the layout (``record_words(world)``). + #: The group size and the row capacity, which fix the layout (``record_words``). self.world = int(mapping.tp_size) - self.record_words = record_words(self.world) + self.max_rows = int(max_rows) + self.wide = self.max_rows > MAX_ROWS + self.record_words = record_words(self.world, self.max_rows) self.guard_ns = int(guard_ns) self.guard_break = bool(guard_break) self.device = torch.device("cuda", _mnnvl_device_index(mapping)) @@ -1040,7 +1248,7 @@ def __init__( error: Optional[BaseException] = None try: use_fabric = mapping.is_multi_node() - nbytes = workspace_bytes(self.world) + nbytes = workspace_bytes(self.world, self.max_rows) self._handle = _make_mnnvl_mcast_buffer(self._comm, nbytes, mapping, use_fabric) words = nbytes // 4 self._receive = self._handle.get_uc_buffer(self.rank, (words,), torch.int32, 0) @@ -1048,6 +1256,12 @@ def __init__( self._records = self._mc.view(BUFFERS, self.world, self.record_words) with torch.cuda.device(self.device): self._receive.fill_(_SENTINEL_I32) + # The wide layout's half-window census words (rank-local, zero between calls). + self._census = ( + torch.zeros(BUFFERS * CENSUS_WORDS, dtype=torch.int32, device=self.device) + if self.wide + else None + ) self.control = torch.zeros( CONTROL_WORDS, dtype=torch.int32, pin_memory=prefer_pinned() ) @@ -1074,7 +1288,7 @@ def slot(self, buffer: int, rows: int) -> torch.Tensor: Write-only: a store through it reaches every rank's copy; nothing may read it. """ - if not 1 <= rows <= MAX_ROWS: + if not 1 <= rows <= self.max_rows: raise ValueError(f"no slot for {rows} rows") record = self._record(buffer) return ( @@ -1088,13 +1302,22 @@ def boxes(self, buffer: int, rows: int) -> torch.Tensor: Multicast and write-only. """ - if not 1 <= rows <= MAX_ROWS: - raise ValueError(f"no boxes for {rows} rows") + if self.wide or not 1 <= rows <= MAX_ROWS: + raise ValueError(f"no boxes for {rows} rows on a {self.max_rows}-row layout") copies = (rows + 1) // 2 record = self._record(buffer) start, box = pool_slot_words(self.world), box_words(self.world) return record[start : start + copies * box].view(torch.float32).view(copies, box) + def block(self, buffer: int, rows: int) -> torch.Tensor: + """This rank's ``[rows, L]`` FP32 logit block of ``buffer`` on the wide layout (multicast; write-only).""" + if not self.wide or not 1 <= rows <= self.max_rows: + raise ValueError(f"no logit block for {rows} rows on a {self.max_rows}-row layout") + local = _route.local_experts(self.world) + record = self._record(buffer) + start = pool_slot_words(self.world, self.max_rows) + return record[start : start + rows * local].view(torch.float32).view(rows, local) + def receive(self) -> torch.Tensor: """This rank's copy of the whole transport, ``int32 [3 * world * record_words(world)]`` (tests).""" return self._receive @@ -1104,14 +1327,20 @@ def control_ptr(self) -> int: return self._control_ptr def is_clean(self) -> bool: - """Whether every word of this rank's copy holds the sentinel (synchronizes; tests).""" + """Whether every word of this rank's copy holds the sentinel and every census word + is zero (synchronizes; tests).""" torch.cuda.synchronize(self.device) - return bool((self._receive == _SENTINEL_I32).all().item()) + clean = bool((self._receive == _SENTINEL_I32).all().item()) + if self._census is not None: + clean = clean and not bool(self._census.any().item()) + return clean def reset(self) -> None: """Every word of this rank's copy back to the sentinel (tests; nothing may be in flight).""" torch.cuda.synchronize(self.device) self._receive.fill_(_SENTINEL_I32) + if self._census is not None: + self._census.zero_() torch.cuda.synchronize(self.device) def guard_fired(self, sync: bool = True) -> bool: @@ -1207,6 +1436,8 @@ def kernel(self, key: Variant): self.guard_ns, self.guard_break, self.world, + self.max_rows, + self._census.data_ptr() if self.wide and place != PLACE_NONE else 0, ) # A dry launch loads the module now, not inside a capture. with torch.cuda.device(self.device): @@ -1278,12 +1509,13 @@ def _launch( _SHARED: dict = {} -def shared_workspace(mapping: Mapping) -> PoolDirectWorkspace: - """The one direct pool workspace of this group, built on first use (collective).""" - workspace = _SHARED.get(mapping) +def shared_workspace(mapping: Mapping, max_rows: int = MAX_ROWS) -> PoolDirectWorkspace: + """The direct pool workspace of this group and row capacity, built on first use (collective).""" + key = (mapping, int(max_rows)) + workspace = _SHARED.get(key) if workspace is None: - workspace = PoolDirectWorkspace(mapping) - _SHARED[mapping] = workspace + workspace = PoolDirectWorkspace(mapping, max_rows=max_rows) + _SHARED[key] = workspace return workspace @@ -1305,6 +1537,68 @@ class PoolDirectRoute(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell( + why="DFlash2 K=7 verify: 8 rows, balanced placement, buffer 0", + spec=dict(rows=8, place=_route.PLACE_BALANCE, buffer=0), + ), + Cell( + why="8 rows over three ring windows, the second buffer", + spec=dict(rows=8, place=_route.PLACE_BALANCE3, buffer=1), + ), + Cell( + why="the third buffer (the last call of an odd count)", + spec=dict(rows=8, place=_route.PLACE_BALANCE, buffer=2), + ), + Cell( + why="8 rows, routing only (no placement)", spec=dict(rows=8, place=PLACE_NONE, buffer=0) + ), + Cell( + why="1 row: one CTA, its second row idle, one logit box", + spec=dict(rows=1, place=_route.PLACE_BALANCE, buffer=1), + ), + Cell( + why="3 rows: the second CTA carries one row", + spec=dict(rows=3, place=_route.PLACE_BALANCE, buffer=0), + ), + Cell( + why="8 rows, owner-only placement", + spec=dict(rows=8, place=_route.PLACE_OWNER, buffer=1), + ), + Cell( + why="8 rows, static placement", + spec=dict(rows=8, place=_route.PLACE_STATIC, buffer=0), + ), + Cell( + why="tep4 wide layout: two requests' K=7 verify (16 rows), routing only", + spec=dict(rows=16, place=PLACE_NONE, buffer=0, wide=True), + ), + Cell( + why="tep4 wide layout: three requests (24 rows), the second buffer", + spec=dict(rows=24, place=PLACE_NONE, buffer=1, wide=True), + ), + Cell( + why="tep4 wide layout: four requests (32 rows), the third buffer", + spec=dict(rows=32, place=PLACE_NONE, buffer=2, wide=True), + ), + Cell( + why="tep4 wide layout: an odd row count (19 rows), the last CTA one row", + spec=dict(rows=19, place=PLACE_NONE, buffer=0, wide=True), + ), + Cell( + why="tep4 wide layout: 16 rows on the half window, the grid's census", + spec=dict(rows=16, place=PLACE_HALF, buffer=1, wide=True), + ), + Cell( + why="tep4 wide layout: 32 rows on the half window (16 CTAs at the census)", + spec=dict(rows=32, place=PLACE_HALF, buffer=0, wide=True), + ), + Cell( + why="tep4 wide layout: 19 rows on the half window, the last CTA one row", + spec=dict(rows=19, place=PLACE_HALF, buffer=2, wide=True), + ), + ) + note = """ Certified against `pool_gather_route` on the same slices and logits bit for bit (the test publishes with the production split-K pool GEMM and the router's @@ -1421,9 +1715,15 @@ def is_valid( stamps: Optional[torch.Tensor] = None, ) -> None: """Raw vector accesses at fixed strides on the operands' device; the transport's - buffer and row count must be the ones the producers published into.""" - if not 1 <= rows <= MAX_ROWS: - raise ValueError(f"rows must be in 1..{MAX_ROWS}, got {rows}") + buffer and row count must be the ones the producers published into. The wide + layout routes only (``PLACE_NONE``, no early ids).""" + if not 1 <= rows <= workspace.max_rows: + raise ValueError(f"rows must be in 1..{workspace.max_rows}, got {rows}") + if workspace.wide and (place not in (PLACE_NONE, PLACE_HALF) or early is not None): + raise ValueError( + "the wide layout routes only or places the half window (PLACE_NONE or PLACE_HALF), " + "without early ids" + ) if buffer not in range(BUFFERS): raise ValueError(f"buffer must be in 0..{BUFFERS - 1}, got {buffer}") if not enable_pdl and early is not None: diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/unpool_direct_boundary.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/unpool_direct_boundary.py index f6088828455f..c41d6ac9a615 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/unpool_direct_boundary.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/comm/unpool_direct_boundary.py @@ -56,6 +56,12 @@ keeps polling): a call sequence that differs between ranks, or a peer that stopped, shows up there, and ``start_guard_watch`` logs it from the host. +Wide transport (TP4, ``max_rows=WIDE_MAX_ROWS``): the decode steps of two to four +requests (16..32 rows, ``wide_prologue``) take a separate workspace whose slots hold 32 +rows (``3 buffers x 4 source ranks x 32 rows x 2048``), with its own buffer sequence; +the kernel is the gather boundary's at those row counts (a cluster of 4 CTAs, one per +peer slice) with the slot stride of 32 rows. The 8-row layout is unchanged. + PDL memory order: before ``griddepcontrol.wait`` the kernel touches shared memory and the cluster barrier only (and, with ``trigger="entry"``, releases its dependents first). After it: the coefficients and weights (immutable), @@ -79,13 +85,13 @@ import torch from cuda.bindings import runtime as _cudart -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from . import unpool_gather_boundary as _f1 -from .projection_inbox import HAS_CUTE_DSL, MAX_ROWS, WORLDS +from .projection_inbox import HAS_CUTE_DSL, MAX_ROWS, WIDE_MAX_ROWS, WORLDS #: Width of the residual stream; one rank's unpool slice is ``local_columns(TP)`` of it. HIDDEN = _f1.HIDDEN @@ -101,6 +107,21 @@ #: Bytes of the whole transport: ``BUFFERS`` x TP slots of ``MAX_ROWS`` rows of a rank's #: ``HIDDEN / TP`` columns, the same at every TP size. WORKSPACE_BYTES = BUFFERS * MAX_ROWS * HIDDEN * 2 + + +def workspace_bytes(max_rows: int = MAX_ROWS) -> int: + """Bytes of the whole transport for slots of ``max_rows`` rows (``WORKSPACE_BYTES`` at 8).""" + return BUFFERS * max_rows * HIDDEN * 2 + + +def check_layout(world: int, max_rows: int) -> None: + """Raise unless ``(world, max_rows)`` is a served layout: 8 rows at TP4 or TP8, 32 rows at TP4.""" + if world not in WORLDS or max_rows not in (MAX_ROWS, WIDE_MAX_ROWS): + raise ValueError(f"no direct unpool layout for TP{world} and {max_rows} rows") + if max_rows == WIDE_MAX_ROWS and world != 4: + raise ValueError("the wide direct unpool layout is TP4 only") + + #: Per-rank control words, in pinned host memory the kernel stores to at system scope: #: ``GUARD_WORD`` is set by a poll that outlived the guard. CONTROL_WORDS = 4 @@ -212,8 +233,14 @@ def __init__( guard_ns: int, guard_break: bool, world: int = 8, + max_rows: int = MAX_ROWS, ): - super().__init__(m, gain, post_eps, pre_eps, norm_fp32, 0, trigger, world) + check_layout(world, max_rows) + super().__init__(m, gain, post_eps, pre_eps, norm_fp32, 0, trigger, world, max_rows) + if not 1 <= m <= max_rows: + raise ValueError(f"{m} rows on a {max_rows}-row layout") + # The slot stride in rows: the workspace's layout. + self.max_rows = int(max_rows) self.capture = bool(capture) self.guard_ns = int(guard_ns) self.guard_break = bool(guard_break) @@ -325,7 +352,7 @@ def kernel( box = ( receive.toint() + ( - ((Int64(parity) * self.world + Int64(peer)) * MAX_ROWS + Int64(row)) + ((Int64(parity) * self.world + Int64(peer)) * self.max_rows + Int64(row)) * self.row_words + Int64(word) ) @@ -459,6 +486,7 @@ def _compile( guard_ns: int, guard_break: bool, world: int = 8, + max_rows: int = MAX_ROWS, ): kernel = _UnpoolDirectBoundaryKernel( rows, @@ -471,6 +499,7 @@ def _compile( guard_ns, guard_break, world, + max_rows, ) pointers = [ make_ptr(dtype, 16, cute.AddressSpace.gmem, assumed_align=16) @@ -560,7 +589,8 @@ class UnpoolDirectWorkspace: ``slot(buffer, rows)`` is this rank's ``[rows, 8192 / TP]`` BF16 slot of buffer ``buffer`` on the multicast address: a GEMM given it as its output publishes - into every rank's copy. ``kernel(variant)`` is the compiled boundary for a + into every rank's copy. ``max_rows`` is the slots' row capacity: 8, or + ``WIDE_MAX_ROWS`` for the wide layout (TP4). ``kernel(variant)`` is the compiled boundary for a call's settings. Construction is collective over the group (fabric allocation and sentinel fill), and so is ``build_variants``; both succeed on every rank or raise on every rank. Nothing compiles inside a CUDA graph @@ -579,6 +609,7 @@ def __init__( *, guard_ns: int = DIRECT_GUARD_NS, guard_break: bool = False, + max_rows: int = MAX_ROWS, ): from tensorrt_llm._torch.distributed.ops import ( _get_mnnvl_workspace_comm, @@ -591,8 +622,11 @@ def __init__( raise ValueError( "the direct unpool transport needs TP4 or TP8 without attention DP or CP" ) + check_layout(int(mapping.tp_size), int(max_rows)) self.rank = mapping.tp_rank self.world = mapping.tp_size + self.max_rows = int(max_rows) + self.wide = self.max_rows > MAX_ROWS self.guard_ns = int(guard_ns) self.guard_break = bool(guard_break) self.device = torch.device("cuda", _mnnvl_device_index(mapping)) @@ -616,14 +650,13 @@ def __init__( error: Optional[BaseException] = None try: use_fabric = mapping.is_multi_node() - self._handle = _make_mnnvl_mcast_buffer( - self._comm, WORKSPACE_BYTES, mapping, use_fabric - ) - words = WORKSPACE_BYTES // 4 + nbytes = workspace_bytes(self.max_rows) + self._handle = _make_mnnvl_mcast_buffer(self._comm, nbytes, mapping, use_fabric) + words = nbytes // 4 self._receive = self._handle.get_uc_buffer(self.rank, (words,), torch.int32, 0) self._mc = self._handle.get_mc_buffer((words,), torch.int32, 0) self._slots = self._mc.view(torch.bfloat16).view( - BUFFERS, self.world, MAX_ROWS, local_columns(self.world) + BUFFERS, self.world, self.max_rows, local_columns(self.world) ) with torch.cuda.device(self.device): self._receive.fill_(_SENTINEL_I32) @@ -653,7 +686,7 @@ def slot(self, buffer: int, rows: int) -> torch.Tensor: Write-only: a store through it reaches every rank's copy; nothing may read it. """ - if buffer not in range(BUFFERS) or not 1 <= rows <= MAX_ROWS: + if buffer not in range(BUFFERS) or not 1 <= rows <= self.max_rows: raise ValueError(f"no slot for buffer {buffer}, {rows} rows") return self._slots[buffer, self.rank, :rows] @@ -774,6 +807,7 @@ def kernel(self, key: Variant): self.guard_ns, self.guard_break, self.world, + self.max_rows, ) # A dry launch loads the module now, not inside a capture. dummy = self._dry.data_ptr() @@ -797,12 +831,13 @@ def publish_copy(slot: torch.Tensor, partial: torch.Tensor) -> None: _SHARED: dict = {} -def shared_workspace(mapping: Mapping) -> UnpoolDirectWorkspace: - """The one direct workspace of this group, built on first use (collective, like the constructor).""" - workspace = _SHARED.get(mapping) +def shared_workspace(mapping: Mapping, max_rows: int = MAX_ROWS) -> UnpoolDirectWorkspace: + """The direct workspace of this group and row capacity, built on first use (collective, like the constructor).""" + key = (mapping, int(max_rows)) + workspace = _SHARED.get(key) if workspace is None: - workspace = UnpoolDirectWorkspace(mapping) - _SHARED[mapping] = workspace + workspace = UnpoolDirectWorkspace(mapping, max_rows=max_rows) + _SHARED[key] = workspace return workspace @@ -824,6 +859,71 @@ class UnpoolDirectBoundary(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell( + why="DFlash2 K=7 verify: 8 rows, the one row count the target publishes at", + spec=dict(rows=8, with_gain=True, norm_fp32_residual=False), + ), + Cell( + why="the second buffer: odd layers publish into buffer 1", + spec=dict(rows=8, with_gain=True, norm_fp32_residual=False, parity=1), + ), + Cell( + why="the third buffer: the last fused layer of an odd count (61 in the model)", + spec=dict(rows=8, with_gain=True, norm_fp32_residual=False, parity=2), + ), + Cell( + why="early release: dependents launched at entry, the next QKV streams its weight meanwhile", + spec=dict(rows=8, with_gain=True, norm_fp32_residual=False, trigger="entry"), + ), + Cell( + why="capture with early release: the DFlash capture from a boundary releasing the next QKV at entry", + spec=dict( + rows=8, with_gain=True, norm_fp32_residual=False, trigger="entry", capture=True + ), + ), + Cell( + why="the FP32-residual pre-QKV norm ordering switch", + spec=dict(rows=8, with_gain=True, norm_fp32_residual=True), + ), + Cell( + why="a weightless first norm (the boundary's other arm)", + spec=dict(rows=8, with_gain=False, norm_fp32_residual=False), + ), + Cell( + why="dependents launched right after the wait", + spec=dict(rows=8, with_gain=True, norm_fp32_residual=False, trigger="after_wait"), + ), + Cell( + why="1 row: the cluster of 16 CTAs, two CTAs per peer slice", + spec=dict(rows=1, with_gain=True, norm_fp32_residual=False), + ), + Cell( + why="3 rows, a short draft between the two cluster shapes", + spec=dict(rows=3, with_gain=True, norm_fp32_residual=False, parity=1), + ), + Cell( + why="tep4 wide layout: two requests' K=7 verify (16 rows), cluster of 4 CTAs", + spec=dict(rows=16, with_gain=True, norm_fp32_residual=False, wide=True), + ), + Cell( + why="tep4 wide layout: 24 rows with the capture and the early release, buffer 1", + spec=dict( + rows=24, + with_gain=True, + norm_fp32_residual=False, + trigger="entry", + capture=True, + parity=1, + wide=True, + ), + ), + Cell( + why="tep4 wide layout: 32 rows (four requests), the third buffer", + spec=dict(rows=32, with_gain=True, norm_fp32_residual=False, parity=2, wide=True), + ), + ) + note = """ Certified against ``unpool_gather_boundary`` on the same slices bit for bit (the test publishes with the production split-K GEMM and with an elementwise copy, through CUDA-graph @@ -941,9 +1041,9 @@ def is_valid( """ rows = published.shape[0] if published.dim() == 2 else -1 columns = local_columns(workspace.world) - if not 1 <= rows <= MAX_ROWS or tuple(published.shape) != (rows, columns): + if not 1 <= rows <= workspace.max_rows or tuple(published.shape) != (rows, columns): raise ValueError( - f"published must be [1..{MAX_ROWS}, {columns}], got {tuple(published.shape)}" + f"published must be [1..{workspace.max_rows}, {columns}], got {tuple(published.shape)}" ) if ( parity not in range(BUFFERS) @@ -988,10 +1088,13 @@ def available() -> bool: return _AVAILABLE -def log_once_enabled(layers: int, kernels: str, buffers: Sequence[int] = ()) -> None: +def log_once_enabled( + layers: int, kernels: str, buffers: Sequence[int] = (), wide_rows: Sequence[int] = () +) -> None: last = f", the last on buffer {buffers[-1]}" if buffers else "" + wide = f", and {'/'.join(map(str, wide_rows))}-row ones on the wide layout" if wide_rows else "" logger.info( f"[musespark_open] unpool_direct: the unpool GEMM publishes into the post-FFN boundary's direct transport on " - f"{layers} layers (8-row decode graphs; buffers alternate per layer{last}; unpool kernel " + f"{layers} layers (8-row decode graphs{wide}; buffers alternate per layer{last}; unpool kernel " f"{kernels}); eager and other row counts keep the gather boundary's tagged inbox" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_allreduce_transport.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_allreduce_transport.py index 2ac6cb991d24..42298bece008 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_allreduce_transport.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/finalize_allreduce_transport.py @@ -52,10 +52,11 @@ for its peers. Protocol (one-shot Lamport; the slot layout is this workspace's own). Each rank -owns a receive area of ``2 buffers x 8 source ranks x 16 tokens x 4096`` BF16 -filled with the sentinel word ``0x80000000``: a 32-bit word that holds it has -not arrived. Published words never equal it (``-0`` halves are cleared). In -one call a slot ``(buffer, source, token, column vector)`` has one writer, the +owns a receive area of ``2 buffers x ranks source ranks x max_tokens tokens x +4096`` BF16 (``max_tokens`` is the workspace's layout, ``MAX_TOKENS`` or +``WIDE_MAX_TOKENS``) filled with the sentinel word ``0x80000000``: a 32-bit +word that holds it has not arrived. Published words never equal it (``-0`` +halves are cleared). In one call a slot ``(buffer, source, token, column vector)`` has one writer, the source rank's thread with that token and column, and one reader, this rank's thread with that token and column, which writes the sentinel back once it has the value -- before its grid completes, hence before this rank publishes its @@ -70,7 +71,10 @@ two calls waits on its predecessor before completing) and before its own CTA arrived, so it reads ``GRID_STEP * n + k`` with ``0 <= k < GRID_STEP``. The grid is at most ``GRID_STEP`` CTAs. ``2**32`` is a multiple of -``2 * GRID_STEP``, so the parity survives the counter's wrap. +``2 * GRID_STEP``, so the parity survives the counter's wrap. ``GRID_STEP`` is +the layout's (``grid_step``): 128 for ``MAX_TOKENS``, 256 for +``WIDE_MAX_TOKENS`` (a 32-token call is 256 CTAs of 64 threads and 2 KiB of +shared memory, resident at once on the devices served). Progress: every CTA publishes before it polls, and a publication waits for no read, so every slot is written; a stale slot is read again every round with a @@ -101,7 +105,15 @@ ``FinalizeTransportWorkspace.finalize`` has the built-in workspace's signature and input contract (its validation is the built-in's): calls of 1..16 tokens -run this kernel, larger ones the built-in workspace it wraps. +run this kernel, larger ones the built-in workspace it wraps. At TP4 a +workspace can also carry a second one laid out for ``WIDE_MAX_TOKENS`` (its own +receive area, arrival counter and kernel build; ``serve_wide_calls``), which +then runs the calls of 17..32 tokens; the 1..16-token calls keep the first +workspace's kernel, memory and counter. The built-in TP4 kernel runs one CTA +per token at every token count up to its 128, so the TP4 arithmetic above is +its arithmetic at 17..32 tokens too; at TP8 the built-in switches to one CTA +per token above 16 tokens, which the cluster-8 order is not, so the wide +layout is TP4-only. """ from __future__ import annotations @@ -112,7 +124,7 @@ import torch -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -124,6 +136,10 @@ TOP_K = 8 #: Tokens per call this kernel serves: the built-in's cluster-8 range at TP8. MAX_TOKENS = 16 +#: Tokens per call of the wide layout (four requests at K = 7), TP4 only. +WIDE_MAX_TOKENS = 32 +#: The layouts a workspace can have, and the group sizes each serves. +LAYOUT_RANK_COUNTS = {MAX_TOKENS: (4, 8), WIDE_MAX_TOKENS: (4,)} #: CTAs per token (one cluster), threads per CTA, BF16 values per thread. CLUSTER = 8 VECTOR = 8 @@ -137,9 +153,14 @@ SLAB_BYTES = MAX_TOKENS * HIDDEN * 2 -def workspace_bytes(ranks: int) -> int: - """Bytes of one rank's receive area for a group of ``ranks``.""" - return BUFFERS * ranks * SLAB_BYTES +def slab_bytes(max_tokens: int = MAX_TOKENS) -> int: + """Bytes of one source rank's slab of one buffer in a ``max_tokens`` layout.""" + return max_tokens * HIDDEN * 2 + + +def workspace_bytes(ranks: int, max_tokens: int = MAX_TOKENS) -> int: + """Bytes of one rank's receive area for a group of ``ranks`` in a ``max_tokens`` layout.""" + return BUFFERS * ranks * slab_bytes(max_tokens) #: Control words: 0 the buffer index and 1 the counter of the built-in-style @@ -149,6 +170,19 @@ def workspace_bytes(ranks: int) -> int: GRID_STEP = 128 GRID_SHIFT = 7 assert GRID_STEP == 1 << GRID_SHIFT and GRID_STEP >= MAX_TOKENS * CLUSTER + + +def grid_shift(max_tokens: int = MAX_TOKENS) -> int: + """``log2`` of what one call adds to the ``red`` counter in a ``max_tokens`` layout. + + The smallest power of two of at least the layout's largest grid + (``max_tokens * CLUSTER`` CTAs): ``GRID_SHIFT`` for ``MAX_TOKENS``. + """ + return (max_tokens * CLUSTER - 1).bit_length() + + +assert grid_shift(MAX_TOKENS) == GRID_SHIFT + #: Bytes one ``push`` value occupies in a receiving CTA (one 16-byte ``st.async``). _PUSH_BYTES = 16 @@ -379,6 +413,8 @@ class LamportExchange: ranks: int multicast: bool + #: One source rank's slab of one buffer; a wider layout sets its own. + slab_bytes: int = SLAB_BYTES def slot_offset(self, token, column): """Byte offset of ``(token, column vector)`` inside one source slab.""" @@ -386,7 +422,7 @@ def slot_offset(self, token, column): def slab_offset(self, buffer, source): """Byte offset of ``source``'s slab of ``buffer`` in a receive area.""" - return (Int64(buffer) * self.ranks + Int64(source)) * SLAB_BYTES + return (Int64(buffer) * self.ranks + Int64(source)) * self.slab_bytes def publish(self, own, buffer, slot, rank, mcast, peers): """Store this rank's four words into slot ``rank`` of every rank's buffer.""" @@ -437,7 +473,7 @@ def poll( while stale != Uint32(0): polled = stale for p in cutlass.range_constexpr(self.ranks): - address = receive + (Int64(buffer) * self.ranks + p) * SLAB_BYTES + slot + address = receive + (Int64(buffer) * self.ranks + p) * self.slab_bytes + slot v0, v1, v2, v3 = _poll_x4( address, (polled >> Uint32(p)) & Uint32(1), @@ -459,7 +495,7 @@ def poll( def retire(self, buffer, slot, rank, receive): """Write the sentinel back into every peer slot this thread consumed.""" for p in range(self.ranks): - address = receive + (Int64(buffer) * self.ranks + p) * SLAB_BYTES + slot + address = receive + (Int64(buffer) * self.ranks + p) * self.slab_bytes + slot _reset_x4(address, Uint32(rank != p)) class _FinalizeTransportKernel(LamportExchange): @@ -469,12 +505,24 @@ class _FinalizeTransportKernel(LamportExchange): each thread owns one 16-byte column vector of that slice for all 8 slots -- the built-in kernel's map. ``ranks`` is the TP size; at TP4 the slot sums are added flat (``flat_sum``, the built-in cluster-1 order). + ``max_tokens`` is the receive-area layout (the slab per source rank and + the arrival counter's step); every statement on a token's values is the + same in every layout. """ - def __init__(self, transport: Transport, pdl: bool, ranks: int = 8): + def __init__( + self, transport: Transport, pdl: bool, ranks: int = 8, max_tokens: int = MAX_TOKENS + ): if ranks not in RANK_COUNTS: raise ValueError(f"the finalize transport serves {RANK_COUNTS} ranks, got {ranks}") + if ranks not in LAYOUT_RANK_COUNTS.get(max_tokens, ()): + raise ValueError( + f"the finalize transport has no {max_tokens}-token layout at {ranks} ranks" + ) self.ranks = ranks + self.slab_bytes = slab_bytes(max_tokens) + self.grid_shift = grid_shift(max_tokens) + self.grid_step = 1 << self.grid_shift # The built-in TP4 kernel runs one CTA per token and adds its warps' sums # from +0 in warp order; TP8's cluster of 8 adds per-CTA totals. self.flat_sum = ranks == 4 @@ -658,10 +706,10 @@ def kernel( _store_control(control_address, current ^ Uint32(1)) else: counter = _load_control(control_address + 8) - snapshot[0] = Int32((counter >> Uint32(GRID_SHIFT)) & Uint32(1)) + snapshot[0] = Int32((counter >> Uint32(self.grid_shift)) & Uint32(1)) amount = Uint32(1) if block == 0: - amount = Uint32(GRID_STEP + 1) - Uint32(blocks) + amount = Uint32(self.grid_step + 1) - Uint32(blocks) _red_add(control_address + 8, amount) gains = [] @@ -816,8 +864,14 @@ def kernel( _POINTER_ALIGNMENTS = (16, 4, 4) + (16,) * (len(_POINTER_DTYPES) - 3) @functools.lru_cache(maxsize=32) - def _compiled_kernel(device_index: int, transport: Transport, pdl: bool, ranks: int = 8): - kernel = _FinalizeTransportKernel(transport, pdl, ranks) + def _compiled_kernel( + device_index: int, + transport: Transport, + pdl: bool, + ranks: int = 8, + max_tokens: int = MAX_TOKENS, + ): + kernel = _FinalizeTransportKernel(transport, pdl, ranks, max_tokens) with torch.cuda.device(device_index): stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) return cute.compile( @@ -845,19 +899,29 @@ class FinalizeTransportWorkspace: """This transport's memory and state for one TP4 or TP8 group, in front of a built-in workspace. ``finalize`` has the built-in ``FinalizeIpcWorkspace.finalize`` signature and - input contract: calls of 1..16 tokens run this module's kernel, every other - call the wrapped built-in workspace (whose own protocol state is separate, - so the two interleave freely; every rank routes a call alike because the - choice depends on ``num_tokens`` only). Construction is collective over the - group (multicast allocation, sentinel fill, kernel compilation and a dry - launch) and either succeeds on every rank or raises on every rank. + input contract: calls of 1..``max_tokens`` tokens run this module's kernel, + calls of the wide workspace's range (``serve_wide_calls``, TP4) run its + kernel, and every other call the wrapped built-in workspace (each has its + own protocol state, so they interleave freely; every rank routes a call + alike because the choice depends on ``num_tokens`` only). Construction is + collective over the group (multicast allocation, sentinel fill, kernel + compilation and a dry launch) and either succeeds on every rank or raises + on every rank. The wrapped workspace is the TP8 fabric one (two trays) or the TP4 IPC one (one tray); the multicast transport uses fabric handles only when the group - spans nodes, as the post-attention hop's does. + spans nodes, as the post-attention hop's does. ``max_tokens`` is the + receive-area layout: ``MAX_TOKENS`` at TP8 or TP4, ``WIDE_MAX_TOKENS`` at + TP4 only. """ - def __init__(self, legacy, mapping: Mapping, transport: Transport = DEFAULT): + def __init__( + self, + legacy, + mapping: Mapping, + transport: Transport = DEFAULT, + max_tokens: int = MAX_TOKENS, + ): from tensorrt_llm._torch.distributed.ops import ( _get_mnnvl_workspace_comm, _make_mnnvl_mcast_buffer, @@ -870,10 +934,19 @@ def __init__(self, legacy, mapping: Mapping, transport: Transport = DEFAULT): f"the finalize transport serves TP{' or TP'.join(map(str, RANK_COUNTS))}, " f"got TP{mapping.tp_size}" ) + if mapping.tp_size not in LAYOUT_RANK_COUNTS.get(max_tokens, ()): + raise ValueError( + f"the finalize transport has no {max_tokens}-token layout at TP{mapping.tp_size}" + ) self.legacy = legacy self.transport = transport self.ranks = mapping.tp_size self.rank = mapping.tp_rank + self.max_tokens = int(max_tokens) + self.slab_bytes = slab_bytes(self.max_tokens) + self.grid_step = 1 << grid_shift(self.max_tokens) + #: The workspace of the calls above ``max_tokens`` (``serve_wide_calls``), or None. + self.wide: Optional["FinalizeTransportWorkspace"] = None self.pdl = _pdl_enabled() self.device = torch.device("cuda", _mnnvl_device_index(mapping)) self._comm = _get_mnnvl_workspace_comm(mapping) @@ -891,9 +964,16 @@ def __init__(self, legacy, mapping: Mapping, transport: Transport = DEFAULT): raise RuntimeError("the wrapped TP8 workspace is not the fabric one") if self.ranks == 4 and (legacy._use_mnnvl or mapping.is_multi_node()): raise RuntimeError("the wrapped TP4 workspace is not the single-tray IPC one") + if _legacy._cluster_for_tokens(self.max_tokens, self.ranks) != ( + CLUSTER if self.ranks == 8 else 1 + ): + raise RuntimeError( + f"the built-in TP{self.ranks} kernel does not split a {self.max_tokens}-token " + "call the way this kernel's order assumes" + ) if torch.cuda.get_device_capability(self.device) not in ((10, 3), (10, 7)): raise RuntimeError("the finalize transport runs on SM103 or SM107 only") - size = workspace_bytes(self.ranks) + size = workspace_bytes(self.ranks, self.max_tokens) self._handle = _make_mnnvl_mcast_buffer( self._comm, size, mapping, mapping.is_multi_node() ) @@ -925,7 +1005,9 @@ def __init__(self, legacy, mapping: Mapping, transport: Transport = DEFAULT): ) error = None try: - self._compiled = _compiled_kernel(self.device.index, transport, self.pdl, self.ranks) + self._compiled = _compiled_kernel( + self.device.index, transport, self.pdl, self.ranks, self.max_tokens + ) # A dry launch loads the module now, not inside a capture. dummy = self.control.data_ptr() self._launch(dummy, dummy, dummy, dummy, dummy, 0, HIDDEN, 1, 1.0, 1) @@ -955,14 +1037,41 @@ def _launch(self, gemm, indices, weights, norm, out, rows, row_stride, tokens, e cuda.CUstream(torch.cuda.current_stream(self.device).cuda_stream), ) - @staticmethod - def takes(num_tokens) -> bool: - """Whether a call of ``num_tokens`` runs this kernel (else the built-in workspace).""" + def takes(self, num_tokens) -> bool: + """Whether a call of ``num_tokens`` runs this workspace's own kernel.""" return ( isinstance(num_tokens, int) and not isinstance(num_tokens, bool) - and 1 <= num_tokens <= MAX_TOKENS + and 1 <= num_tokens <= self.max_tokens + ) + + def serves(self, num_tokens) -> bool: + """Whether a call of ``num_tokens`` runs this kernel or the wide workspace's, not the built-in.""" + return self.takes(num_tokens) or (self.wide is not None and self.wide.takes(num_tokens)) + + def serve_wide_calls(self, mapping: Mapping) -> "FinalizeTransportWorkspace": + """Run calls of ``max_tokens + 1..WIDE_MAX_TOKENS`` tokens on a second workspace (collective). + + The second workspace has the ``WIDE_MAX_TOKENS`` layout, the same transport + and the same wrapped built-in workspace, and its own receive area, arrival + counter and kernel build; this workspace's calls of 1..``max_tokens`` + tokens keep its own. TP4 only (``LAYOUT_RANK_COUNTS``). Returns the + second workspace. + + Raises: + ValueError: at a group size without the wide layout, or when this + workspace already serves its wide calls or already takes them. + RuntimeError: if the second workspace could not be built on some rank + (on every rank). + """ + if self.wide is not None: + raise ValueError("the finalize transport already serves its wide calls") + if self.max_tokens >= WIDE_MAX_TOKENS: + raise ValueError(f"the finalize transport already takes {self.max_tokens} tokens") + self.wide = FinalizeTransportWorkspace( + self.legacy, mapping, self.transport, WIDE_MAX_TOKENS ) + return self.wide def finalize( self, @@ -973,8 +1082,15 @@ def finalize( eps: float, num_tokens: int, ) -> torch.Tensor: - """``FinalizeIpcWorkspace.finalize``, with 1..16 tokens on this transport.""" + """``FinalizeIpcWorkspace.finalize``, with 1..``max_tokens`` tokens on this transport. + + Calls the wide workspace takes go to it, every other call to the built-in. + """ if not self.takes(num_tokens): + if self.wide is not None and self.wide.takes(num_tokens): + return self.wide.finalize( + gemm2_output, indices, routing_weights, norm_weight, eps, num_tokens + ) return self.legacy.finalize( gemm2_output, indices, routing_weights, norm_weight, eps, num_tokens ) @@ -1001,22 +1117,64 @@ def finalize( class FinalizeAllReduceTransport(OpWrapper): """The fused expert finalize + all-reduce of a ``FinalizeTransportWorkspace``. - Returns the ``[num_tokens, 4096]`` BF16 group sum of every rank's partial: - per token, each valid expert row ``x`` (slot index in ``[0, R)``) scaled by - ``w * rsqrt(mean((x * g)^2) + eps)`` times ``x * g`` with ``g = 1 + norm``, - summed over the slots, rounded to BF16 with ``-0`` cleared, then summed over - the ranks in FP32 in rank order and rounded to BF16. + Returns the ``[num_tokens, 4096]`` BF16 group sum of every rank's partial + (calls of 1..16 tokens, and of 17..32 on a TP4 workspace serving its wide + calls): per token, each valid expert row ``x`` (slot index in ``[0, R)``) + scaled by ``w * rsqrt(mean((x * g)^2) + eps)`` times ``x * g`` with + ``g = 1 + norm``, summed over the slots, rounded to BF16 with ``-0`` + cleared, then summed over the ranks in FP32 in rank order and rounded to BF16. """ ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell(why="no speculation: one token", spec=dict(tokens=1)), + Cell(why="DFlash2 K=7 verify: 8 tokens", spec=dict(tokens=8)), + Cell( + why="the widest call on this kernel (the ring-buddy decode bound)", spec=dict(tokens=16) + ), + Cell( + why="every slot of a token on one rank, none on the others", + spec=dict(tokens=8, placement="concentrated"), + ), + Cell( + why="ring-buddy placement: each slot on its owner or successor", + spec=dict(tokens=8, placement="buddy"), + ), + Cell( + why="exact and negative zeros in the rows, tokens with no valid slot", + spec=dict(tokens=8, zeros=True), + ), + Cell( + why="invalid slots and cancelling rows: the rank-order sum is observable", + spec=dict(tokens=8, invalid=True, cancel=True), + ), + Cell( + why="wide TP4 layout: three requests at K=7, 24 tokens", + spec=dict(tokens=24, placement="buddy", wide=True), + ), + Cell( + why="wide TP4 layout: four requests at K=7, 32 tokens, its widest call", + spec=dict(tokens=32, placement="buddy", wide=True), + ), + Cell( + why="wide TP4 layout: zeros, negative zeros and a token with no valid slot", + spec=dict(tokens=24, zeros=True, wide=True), + ), + Cell( + why="wide TP4 layout: invalid slots and cancelling rows", + spec=dict(tokens=32, invalid=True, cancel=True, wide=True), + ), + ) + note = """ Metadata consumed: ``indices`` (the expert op's ``[tokens, 8]`` row map, read through one function, ``select_row``) and the routing weights. Call site: the built-in MoE's ``workspace.finalize`` once per layer per step, unchanged; the feature puts a ``FinalizeTransportWorkspace`` in that - attribute, which runs this kernel for 1..16 tokens and the built-in kernel - otherwise. + attribute, which runs this kernel for 1..16 tokens (and, with + ``wide_finalize_transport`` at TP4, the same kernel on a 32-token layout + for 17..32 tokens) and the built-in kernel otherwise. Certified bit for bit against the built-in kernel by its own multi-GPU test (8 ranks on two nodes; every cell, every transport variant, tokens 1..16, a @@ -1025,8 +1183,12 @@ class FinalizeAllReduceTransport(OpWrapper): its dependents before it writes their inputs) and against `reference` within a band; the alternation of the two receive buffers by the counter's parity is checked white-box; at TP4 (4 ranks on one tray) the same test - against the built-in TP4 IPC kernel (cluster 1). Not certified: more than - 16 tokens (served by the built-in kernel). + against the built-in TP4 IPC kernel (cluster 1), and there the wide + workspace for every token count 17..32 (the wide cells, graph replays and + bursts mixing both workspaces, a delayed rank, its counter across the wrap, + its buffer parity). Not certified: more than 32 tokens, or more than 16 at + TP8 (served by the built-in kernel). The wide cells (``wide=True``) run at + TP4 only. """ def __call__( @@ -1039,8 +1201,11 @@ def __call__( num_tokens: int, workspace: FinalizeTransportWorkspace, ) -> torch.Tensor: - if not FinalizeTransportWorkspace.takes(num_tokens): - raise ValueError(f"finalize_allreduce_transport serves 1..{MAX_TOKENS} tokens") + if not workspace.serves(num_tokens): + raise ValueError( + f"finalize_allreduce_transport serves 1..{MAX_TOKENS} tokens " + f"(1..{WIDE_MAX_TOKENS} on a TP4 workspace serving its wide calls)" + ) return workspace.finalize( gemm2_output, indices, routing_weights, norm_weight, eps, num_tokens ) @@ -1081,9 +1246,12 @@ def reference( finalize_allreduce_transport = FinalizeAllReduceTransport() -def log_once_enabled(layers: int, transport: Transport) -> None: +def log_once_enabled(layers: int, transport: Transport, wide: bool = False) -> None: + calls = f"1..{MAX_TOKENS} tokens" + if wide: + calls += f" and, on a second receive area, {MAX_TOKENS + 1}..{WIDE_MAX_TOKENS} tokens" logger.info( f"[musespark_open] finalize_transport: fused expert finalize + all-reduce of all {layers} layers on the " f"{transport.label()} transport (publish/arrival/reduction/trigger) for calls of " - f"1..{MAX_TOKENS} tokens; the built-in kernel above that" + f"{calls}; the built-in kernel above that" ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_shard.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_shard.py index 9d80c144ed58..4b8136e127d6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_shard.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/moe/router_shard.py @@ -39,7 +39,7 @@ import triton import triton.language as tl -from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, OpWrapper +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper #: The checkpoint's router: 256 experts over the 8192-wide residual. EXPERTS = 256 @@ -144,6 +144,26 @@ class RouterShard(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell( + why="K=7 verify: 8 rows, the served launch (prepare stream, 16 warps)", + spec=dict(rows=8, num_warps=16), + ), + Cell(why="one row at 16 warps (no speculation)", spec=dict(rows=1, num_warps=16)), + Cell( + why="an odd row count the built-in launches at 1-row tiles (3 rows)", + spec=dict(rows=3, num_warps=16), + ), + Cell( + why="8 rows without the prepare stream: the 4-warp summation order", + spec=dict(rows=8, num_warps=4), + ), + Cell( + why="tep4 K=7 verify: 8 rows, 64 columns per rank (one tray)", + spec=dict(rows=8, num_warps=16, ranks=4), + ), + ) + note = """ The claim the test proves is equality with the built-in kernel's columns ``L r .. L r + L - 1`` at the same warp count, for every rank, every row count @@ -265,6 +285,22 @@ class RouterShardSide(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell( + why="K=7 verify: 8 rows, the served launch (prepare stream, 16 warps)", + spec=dict(rows=8, num_warps=16), + ), + Cell(why="one row at 16 warps (no speculation)", spec=dict(rows=1, num_warps=16)), + Cell( + why="an odd row count the built-in launches at 1-row tiles (3 rows)", + spec=dict(rows=3, num_warps=16), + ), + Cell( + why="8 rows without the prepare stream: the 4-warp summation order", + spec=dict(rows=8, num_warps=4), + ), + ) + note = """ The claim the test proves is equality with ``router_shard`` and with the built-in kernel's columns ``L r .. L r + L - 1`` at the same warp count, for @@ -419,6 +455,22 @@ class RouterShardPublish(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell( + why="K=7 verify: 8 rows, the served launch (16 warps), four boxes", + spec=dict(rows=8, num_warps=16), + ), + Cell(why="one row at 16 warps: one box", spec=dict(rows=1, num_warps=16)), + Cell( + why="an odd row count (3 rows): two boxes, the second CTA one row", + spec=dict(rows=3, num_warps=16), + ), + Cell( + why="8 rows without the prepare stream: the 4-warp summation order", + spec=dict(rows=8, num_warps=4), + ), + ) + note = """ The claim the test proves is that every box holds ``router_shard``'s logits bit for bit (and the built-in kernel's columns ``32r .. 32r + 31``) for every @@ -580,6 +632,17 @@ class RouterShardWide(OpWrapper): ARCHS = frozenset({Arch.SM_103}) + CELLS: tuple[Cell, ...] = ( + Cell( + why="tep4: two requests' K=7 verify, 16 rows, 64 columns per rank", + spec=dict(rows=16, ranks=4), + ), + Cell(why="tep4: three requests, 24 rows (two token tiles)", spec=dict(rows=24, ranks=4)), + Cell(why="tep4: four requests, 32 rows (the wide capacity)", spec=dict(rows=32, ranks=4)), + Cell(why="a row count off the K=7 grid (19 rows)", spec=dict(rows=19, ranks=4)), + Cell(why="tep8 layout: 16 rows, 32 columns per rank", spec=dict(rows=16, ranks=8)), + ) + note = """ The claim the test proves is equality with the built-in router's columns ``L r .. L r + L - 1`` at every row count 16..32, for every rank, at EP4 and EP8, @@ -669,3 +732,189 @@ def is_valid( router_shard_wide = RouterShardWide() + + +#: ``router_shard_rows`` (``wide_prologue``) launches ``_router_shard_side_kernel``'s statements +#: with each program looping over ``PER`` consecutive rows of one expert: the weight row is +#: loaded once per program, and every row's logit is the same ``tl.sum`` over the same +#: ``[1, 8192]`` tile layout (the warp count alone fixes it), so each row's bits are those +#: of ``router_shard`` on that row. Placed after every other kernel of this file. +@triton.jit +def _router_shard_rows_kernel( + hidden_states, + weight, + output, + scale, + TOKENS: tl.constexpr, + PER: tl.constexpr, + EXPERTS: tl.constexpr, + HIDDEN: tl.constexpr, + ROWS: tl.constexpr, +): + first = tl.program_id(0) * PER + expert = tl.program_id(1) * ROWS + tl.arange(0, ROWS) + offsets = tl.arange(0, HIDDEN) + w = tl.load( + weight + expert[:, None] * HIDDEN + offsets[None, :], + expert[:, None] < EXPERTS, + other=0.0, + ) + for k in tl.static_range(PER): + token = first + k + if token < TOKENS: + x = tl.load(hidden_states + token * HIDDEN + offsets).to(tl.float32) + logits = tl.sum(w * x[None, :], axis=1) * scale + tl.store(output + token * EXPERTS + expert, logits, expert < EXPERTS) + + +#: Rows a ``router_shard_rows`` call may carry: two to four requests' K=7 verify steps. +ROWS_MAX_ROWS = 32 +#: Rows per program by row count: the fastest of the bitwise-equal launches on one R200 at +#: 16 warps (every entry writes the same bits). +ROWS_PER_PROGRAM = {16: 1, 24: 4, 32: 2} + + +def rows_per_program(tokens: int) -> int: + """Rows each ``router_shard_rows`` program loops over at ``tokens`` rows (2 off the K=7 grid).""" + return ROWS_PER_PROGRAM.get(tokens, 2) + + +class RouterShardRows(OpWrapper): + """The rank's ``L`` router logit columns of up to 32 rows, each row's one-request bits. + + ``router_shard_side``'s statements -- the built-in small-batch router's -- over every row + of a 2..4-request decode step: each logit is the one ``router_shard`` computes for that row + at the same warp count, i.e. the bits the built-in router gives the row when its request + decodes alone (``num_warps = builtin_warps(MAX_ROWS, overlaps)``). The built-in router of a + 16..32-row step instead takes its TF32x3 tensor-core launch, which ``router_shard_wide`` + reproduces; this entry trades those bits for the one-request ones in one launch and no + reduce. A plain launch (no programmatic dependent launch) into the caller's ``out``, for + the router fork's side stream. + + Args: + hidden_states: ``[M, 8192]`` BF16, contiguous, ``1 <= M <= 32``. + weight: the full router weight ``[256, 8192]`` FP32, contiguous. + rank: the expert-parallel rank ``r``; columns ``[L r, L r + L)`` are computed. + scale: the metap multiplier, a Python float (applied in FP32 after the sum). + num_warps: the warp count of the built-in launch at one request's rows (4 or 16). + out: ``[M, L]`` FP32, contiguous, on the input's device; overwritten. + ranks: the expert-parallel group size (4 or 8); ``L = local_experts(ranks)``. + + Returns: + ``out``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="tep4: two requests' K=7 verify, 16 rows, 64 columns, the served 16 warps", + spec=dict(rows=16, num_warps=16, ranks=4), + ), + Cell( + why="tep4: three requests, 24 rows (4 rows per program)", + spec=dict(rows=24, num_warps=16, ranks=4), + ), + Cell( + why="tep4: four requests, 32 rows (the wide capacity, 2 rows per program)", + spec=dict(rows=32, num_warps=16, ranks=4), + ), + Cell( + why="a row count off the K=7 grid (19 rows, a partial last program)", + spec=dict(rows=19, num_warps=16, ranks=4), + ), + Cell( + why="without the prepare stream: the 4-warp summation order", + spec=dict(rows=24, num_warps=4, ranks=4), + ), + Cell(why="tep8 layout: 16 rows, 32 columns", spec=dict(rows=16, num_warps=16, ranks=8)), + ) + + note = """ + The claim the test proves is that every row of every rank's output is ``router_shard``'s + logit for that row bit for bit -- the built-in small-batch router's columns + ``L r .. L r + L - 1`` of that row in a one-request (8-row) call -- at every row count + 1..32, both warp counts, EP4 and EP8, eagerly and through a CUDA-graph replay on a forked + side stream; the cells additionally gate the op against the float64 reference. + """ + + def __call__( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + rank: int, + scale: float, + num_warps: int, + out: torch.Tensor, + ranks: int = 4, + ) -> torch.Tensor: + tokens = hidden_states.shape[0] + local = local_experts(ranks) + first = rank * local + per = rows_per_program(tokens) + _router_shard_rows_kernel[(triton.cdiv(tokens, per), local // TILE_ROWS)]( + hidden_states, + weight[first : first + local], + out, + scale, + TOKENS=tokens, + PER=per, + EXPERTS=local, + HIDDEN=HIDDEN, + ROWS=TILE_ROWS, + num_warps=num_warps, + enable_fp_fusion=False, + launch_pdl=False, + ) + return out + + def reference( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + rank: int, + scale: float, + num_warps: int, + out: torch.Tensor, + ranks: int = 4, + ) -> torch.Tensor: + """``RouterShard.reference``; ``num_warps`` and ``out`` do not change the values.""" + del out + return router_shard.reference(hidden_states, weight, rank, scale, num_warps, ranks=ranks) + + def is_valid( + self, + hidden_states: torch.Tensor, + weight: torch.Tensor, + rank: int, + scale: float, + num_warps: int, + out: torch.Tensor, + ranks: int = 4, + ) -> None: + """Dense row-major operands; 1..32 rows; a built-in warp count.""" + if hidden_states.dtype != torch.bfloat16 or hidden_states.dim() != 2: + raise ValueError(f"hidden_states must be 2-D BF16, got {hidden_states.dtype}") + if not hidden_states.is_contiguous() or hidden_states.shape[1] != HIDDEN: + raise ValueError(f"hidden_states must be contiguous [M, {HIDDEN}]") + rows = hidden_states.shape[0] + if not 1 <= rows <= ROWS_MAX_ROWS: + raise ValueError(f"1..{ROWS_MAX_ROWS} rows, got {rows}") + if weight.dtype != torch.float32 or tuple(weight.shape) != (EXPERTS, HIDDEN): + raise ValueError(f"weight must be FP32 [{EXPERTS}, {HIDDEN}]") + if not weight.is_contiguous(): + raise ValueError("weight must be contiguous") + if ranks not in RANK_COUNTS: + raise ValueError(f"ranks must be one of {RANK_COUNTS}, got {ranks}") + if not 0 <= rank < ranks: + raise ValueError(f"rank must be in [0, {ranks}), got {rank}") + if num_warps not in WARP_COUNTS: + raise ValueError(f"num_warps must be one of {WARP_COUNTS}, got {num_warps}") + local = local_experts(ranks) + if out.dtype != torch.float32 or tuple(out.shape) != (rows, local): + raise ValueError(f"out must be FP32 [{rows}, {local}]") + if not out.is_contiguous() or out.device != hidden_states.device: + raise ValueError("out must be contiguous, on the input's device") + + +router_shard_rows = RouterShardRows() From 4f65cb9994a7f6eaf0d2db439a92a3519c9959d6 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:11:04 +0000 Subject: [PATCH 09/15] [None][perf] modeling_v2 musespark_open tep: single-request post-target chain and a cooperative head-sampler fold Paths for decode steps of one or a few requests on sm107: - drafter K/V side chain: on a dense full-length drafter arena, the context K/V of the step is projected and stored on a side stream forked after the last capture layer, with the drafter's fixed GEMM tactics; the fc projection can fork early, and the commit kernel also writes the drafter forward's integer preamble (context_kv_store, context_kv_commit, context_kv_fork, verify_fold, drafter_block_decode). - head sampler: a cooperative kernel A computes the per-row states of calls of up to 128 rows with the same per-element cap and combine order as the grid kernel, so states and tokens are bitwise the same; wider calls keep the grid kernel. - the target head can serve without the lm_head inbox, and can gather its logits with a multimem all-gather when every rank builds the symmetric buffer. - runtime: the block-offset stager copies only the runs of page indices that changed since the previous step, keyed by value, and resets on storage, scale or offset changes. Each path declines per step when its arena or shape does not apply. Which of them a deployment serves is decided by the lever resolver. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../kernels/context_kv_fork/fork.py | 380 +++++++++++++++++- .../drafter_block_decode/block_decode.py | 80 +++- .../kernels/split_sampler/target_head.py | 202 +++++++++- .../kernels/verify_fold/fold.py | 115 +++++- .../ops/embedding/replicated_gather.py | 50 +++ .../ops/spec/_head_sampler/head_sampler.cu | 188 +++++++++ .../ops/spec/context_kv_commit.py | 160 ++++++++ .../ops/spec/context_kv_store.py | 292 ++++++++++++++ .../ops/spec/head_sampler.py | 30 +- .../runtime.py | 249 +++++++++++- 10 files changed, 1690 insertions(+), 56 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_store.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py index b8d4a6d63f08..017d8da9d02e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/context_kv_fork/fork.py @@ -217,6 +217,17 @@ def __init__(self, worker: DFlashWorker, draft_model: nn.Module, point: str): self._armed: Optional[tuple] = None self._launched: Optional[Launched] = None self._fallback_logged = False + #: The early part of the armed step's side chain (``fire_early``), or None. + self._early = None + self.early_launches = 0 + #: ``drafter_kv_side_chain``: the side chain runs ``_fused_side_chain`` where it applies. + self.side_chain_on = False + #: ``drafter_fc_early_fork``: the fused side chain's local part forks early (``fire_early``). + self.early_fc_on = False + #: ``windowed_commit_parallel``'s MIN_REQUESTS for the windowed commit, or None. + self.parallel_min_requests: Optional[int] = None + #: Reports a step a lever left to the path below it (``dispatch.Plan.declined``), or None. + self.declined: Optional[Callable[..., None]] = None @property def launched(self) -> bool: @@ -230,6 +241,47 @@ def arm(self, spec_metadata, attn_metadata) -> None: return self._armed = (spec_metadata, attn_metadata) + def _decline(self, lever: str, reason: str, num_gens: int) -> None: + if self.declined is not None: + self.declined(lever, reason, f"{num_gens} requests") + + def fire_early(self) -> None: + """The last capture layer has run: launch the fused side chain's local part (``_side_chain_early``). + + Only for an armed step the side chain reproduces (``plan``) and that has not + forked yet. The side stream waits for the main stream's work so far (the capture + rows are complete); the projection gather and the rest fork at the step's fork point + (``fire``), behind another wait for the main stream. + """ + if self._armed is None or self._launched is not None or self._early is not None: + return + if not (self.early_fc_on and self.side_chain_on): + return + spec_metadata, attn_metadata = self._armed + # plan() first: it declines a step whose arena is not built yet (warmup), which the + # reason check below reads. + plan = self.plan(spec_metadata, attn_metadata) + if plan is None: + return + if _side_chain_refusal(self.worker, self.draft_model) is not None: + return + num_gens, k = plan + main = torch.cuda.current_stream() + if main == self.stream: + return + self.stream.wait_stream(main) + with torch.cuda.stream(self.stream): + early = _side_chain_early(self.worker, self.draft_model, spec_metadata, num_gens, k) + if early is None: + self._decline( + "drafter_fc_early_fork", + "the drafter's projection is not the sharded fc on this step", + num_gens, + ) + return + self._early = (early, num_gens, k) + self.early_launches += 1 + def fire(self, point: str) -> None: """A fork point was reached: launch the armed step's side chain if this is its point.""" if self._armed is None or self._launched is not None: @@ -258,6 +310,7 @@ def fire(self, point: str) -> None: self.stream.wait_stream(main) with torch.cuda.stream(self.stream): gen_hs = self.side_chain(spec_metadata, attn_metadata, num_gens, k) + self._early = None self._launched = Launched(main, capture, num_gens, k, point, gen_hs) self.launches[point] += 1 @@ -335,12 +388,77 @@ def commit(self, launched: Launched, slots, num_accepted) -> Optional[torch.Tens worker._max_ctx, column_limit(worker), worker._ctx_keep, + parallel_min_requests=self.parallel_min_requests, ) # The launch advanced and compacted; the worker's rules for the rest. worker._empty_dummy_slot() num_ctx.copy_(worker._drafter_context_lengths(slots, rows)) return num_ctx + def commit_with_preamble( + self, + launched: Launched, + slots, + num_accepted, + query_positions, + seq_lens_cuda, + seq_lens_value: int, + draft_block: int, + ) -> tuple: + """``commit`` whose kernel also writes the drafter preamble (``context_kv_commit_preamble``). + + The dense full-length arena only (the caller checks: no pages, no window), where + ``commit`` is ``context_kv_commit`` and ``_commit_context``; the kernel is both. + Returns ``(num_ctx, i32s)``; every output is allocated on the main stream, as + ``commit``'s, and valid after the join. + """ + from ...ops.spec.context_kv_commit import context_kv_commit_preamble + + worker = self.worker + if ( + worker._ctx_paged + or worker._ctx_block_tables is not None + or worker._ctx_base is not None + ): + raise RuntimeError( + "[musespark_open] context_kv_fork: the preamble commit is the dense full-length arena's; " + "this worker's context is paged or windowed" + ) + main = torch.cuda.current_stream() + if main != launched.main or current_capture_id() != launched.capture: + raise RuntimeError( + "[musespark_open] context_kv_fork: the side chain was forked from another stream or " + "capture state than the one committing it" + ) + device = worker._ctx_len.device + gens = launched.num_gens + num_ctx = torch.empty(gens, dtype=torch.int64, device=device) + i32s = ( + torch.empty(gens, dtype=torch.int32, device=device), + torch.empty(gens, dtype=torch.int32, device=device), + torch.empty(gens, dtype=torch.int32, device=device), + torch.empty(gens * draft_block, dtype=torch.int32, device=device), + ) + self.stream.wait_stream(main) + with torch.cuda.stream(self.stream): + context_kv_commit_preamble( + worker._ctx_k_buf, + worker._ctx_v_buf, + worker._ctx_len, + slots, + num_accepted, + num_ctx, + launched.k + 1, + worker._max_ctx, + query_positions, + seq_lens_cuda, + seq_lens_value, + draft_block, + i32s, + worker._dummy_slot, + ) + return num_ctx, i32s + def settle(self) -> None: """Forget an armed step and join a side chain a failed step left behind. @@ -351,6 +469,7 @@ def settle(self) -> None: """ launched, self._launched = self._launched, None self._armed = None + self._early = None if launched is not None and current_capture_id() == launched.capture: torch.cuda.current_stream().wait_stream(self.stream) @@ -404,7 +523,21 @@ def side_chain(self, spec_metadata, attn_metadata, num_gens: int, k: int) -> tor table decode over the KV cache manager's pages. Its first launch is a torch kernel, so the fork is a plain dependency on the main stream's last kernel. """ - worker = self.worker + worker, drafter = self.worker, self.draft_model + refusal = _side_chain_refusal(worker, drafter) if self.side_chain_on else None + if self.side_chain_on and refusal is None: + # The dense full-length arena only (``_side_chain_refusal``): no page table to decode. + early = None + if self._early is not None: + early, early_gens, early_k = self._early + if (early_gens, early_k) != (num_gens, k) or early.capture != current_capture_id(): + raise RuntimeError( + "[musespark_open] context_kv_fork: the early side-chain part was launched for " + "another step or capture state than the one forking now" + ) + return _fused_side_chain(worker, drafter, spec_metadata, num_gens, k, early) + if refusal is not None: + self._decline("drafter_kv_side_chain", refusal, num_gens) DFlashWorker._refresh_ctx_block_tables(worker, attn_metadata, num_gens) slots = worker._batch_to_slot[:num_gens] # The built-in's get_hidden_states(num_tokens)[0 : rows]: the same view. @@ -417,6 +550,232 @@ def side_chain(self, spec_metadata, attn_metadata, num_gens: int, k: int) -> tor return gen_hs +# -------------------------------------------------------------------------- +# The side chain with its integer prologue and its arena store fused +# (``drafter_kv_side_chain``, ``ContextKvFork.side_chain_on``) +# -------------------------------------------------------------------------- + + +def _side_chain_refusal(worker, drafter) -> Optional[str]: + """Why ``_fused_side_chain`` does not write out this drafter's ``precompute_context_kv``, or None. + + The store it writes out is the dense full-length arena's (``context_kv_scatter`` at + columns ``min(ctx_len + r, max_ctx - 1)``): over the KV cache manager's pages or a + windowed arena the built-in ``_store_step_context`` runs (``ContextKvFork.side_chain``). + """ + from tensorrt_llm._torch.models import modeling_dflash + + if worker._ctx_paged or worker._ctx_block_tables is not None: + return "the context is in the KV cache manager's pages" + if worker._ctx_base is not None: + return "the context arena is windowed" + if modeling_dflash._flashinfer_rope is None: + return "no flashinfer RoPE" + if drafter._fused_kv_weight is None: + drafter._build_fused_kv_buffers() + if getattr(drafter, "_input_ln_eps", None) is not None: + return "the drafter normalizes the projected rows before its K/V GEMM" + if drafter._fused_kv_weight.dtype != worker._ctx_k_buf.dtype: + return "the K/V GEMM weight is not the arena dtype" + if drafter._k_norm_stacked is None: + return "no stacked K norm" + if worker._ctx_k_buf.dim() != 5 or worker._ctx_k_buf.stride(4) != 1: + return "the arena is not [slots, layers, capacity, kv_heads, head_dim]" + if _kv_gemm(drafter) is None: + return "precompute_context_kv is replaced by something other than drafter_fixed_gemms" + return None + + +def _kv_gemm(drafter): + """The fused K/V GEMM the drafter's ``precompute_context_kv`` runs, as ``f(x) -> kv_flat``, or None. + + The built-in method (not replaced on the instance): ``F.linear`` with the fused bias; + ``drafter_fixed_gemms``'s ``fixed_tactic_precompute_context_kv`` (bias-free): its + fixed-tactic ``_gemm``. Anything else: None (the stock side chain runs). + """ + from ... import spec_worker as _drafter_worker + + bound = drafter.__dict__.get("precompute_context_kv") + weight, bias = drafter._fused_kv_weight, drafter._fused_kv_bias + if bound is None: + if type(drafter).precompute_context_kv is not DFlashForCausalLM.precompute_context_kv: + return None + return lambda x: torch.nn.functional.linear(x, weight, bias) + if _bound_to(bound, _drafter_worker.fixed_tactic_precompute_context_kv): + if bias is not None: + return lambda x: torch.nn.functional.linear(x, weight, bias) + # The same tactic table as fixed_tactic_precompute_context_kv's call (the drafter's TP size). + pdl = drafter._fixed_gemm_pdl + tactics = getattr(drafter, "_fixed_gemm_tactics", None) + return lambda x: _drafter_worker._gemm(x, weight, pdl, tactics) + return None + + +@dataclass +class _SideChainPrologue: + """The side chain's first part: its prep kernel's outputs and the captured rows.""" + + slots: torch.Tensor + positions: torch.Tensor + col: torch.Tensor + pos_i32: torch.Tensor + gen_hs: torch.Tensor + + +def _side_chain_prologue( + worker, drafter, spec_metadata, num_gens: int, k: int +) -> _SideChainPrologue: + """``context_kv_prep`` (positions, clamped columns, per-layer int32 RoPE positions) and the captured rows.""" + from ...ops.spec.context_kv_store import context_kv_prep + + rows_per_req = k + 1 + rows = num_gens * rows_per_req + layers = drafter._num_attn_layers + device = worker._ctx_len.device + slots = worker._batch_to_slot[:num_gens] + positions = torch.empty(rows, dtype=torch.long, device=device) + col = torch.empty(rows, dtype=torch.long, device=device) + pos_i32 = torch.empty(rows * layers, dtype=torch.int32, device=device) + context_kv_prep( + worker._ctx_len, slots, positions, col, pos_i32, rows_per_req, layers, worker._max_ctx + ) + gen_hs = spec_metadata.captured_hidden_states[ + :rows, : spec_metadata.hidden_size * spec_metadata.num_capture_layers + ] + return _SideChainPrologue(slots, positions, col, pos_i32, gen_hs) + + +# ---- The early fc fork (``drafter_fc_early_fork``, ``ContextKvFork.early_fc_on``) ---- +# The fused side chain's local part -- the prep kernel and the drafter's sharded ``fc`` GEMM +# on this rank's rows -- forks right after the last DFlash capture layer +# (``ContextKvFork.fire_early``); the projection gather and everything after it still fork +# at the LM-head point. Nothing in the early part waits on another rank or another kernel. + + +@dataclass +class _EarlyFc: + """The early part's results: the prologue and this rank's ``fc`` block of the projected rows.""" + + prologue: _SideChainPrologue + local: torch.Tensor + gather: object + capture: Optional[int] + + +def _sharded_fc(drafter, x: torch.Tensor): + """``(shard, state)`` when the drafter's ``project_target_hidden`` is the sharded ``fc`` + (``spec_worker.sharded_project_target_hidden``) and that method would take its sharded + branch for ``x``; else None (the whole projection then runs at the LM-head point).""" + from ... import spec_worker as _drafter_worker + + if not _bound_to( + drafter.__dict__.get("project_target_hidden"), + _drafter_worker.sharded_project_target_hidden, + ): + return None + state = drafter.__dict__.get("_fc_shard") + weight = drafter.fc.weight + if ( + state is None + or drafter.fc.bias is not None + or x.dim() != 2 + or x.shape[0] < 1 + or (state.max_rows is not None and x.shape[0] > state.max_rows) + or x.dtype != torch.bfloat16 + or not weight.is_contiguous() + or tuple(weight.shape) != (state.rows * state.world, x.shape[1]) + ): + return None + begin = state.rank * state.rows + return weight.detach()[begin : begin + state.rows], state + + +def _side_chain_early(worker, drafter, spec_metadata, num_gens: int, k: int): + """The side chain's local part, or None when the projection is not the sharded ``fc``. + + ``sharded_project_target_hidden``'s statements up to its gather: ``x = rows.to(weight + dtype)``, ``_gemm(x.contiguous(), shard, pdl)``. + """ + from ... import spec_worker as _drafter_worker + + rows = num_gens * (k + 1) + gen_hs = spec_metadata.captured_hidden_states[ + :rows, : spec_metadata.hidden_size * spec_metadata.num_capture_layers + ] + x = gen_hs.reshape(-1, gen_hs.shape[-1]).to(drafter.fc.weight.dtype) + split = _sharded_fc(drafter, x) + if split is None: + return None + shard, state = split + prologue = _side_chain_prologue(worker, drafter, spec_metadata, num_gens, k) + local = _drafter_worker._gemm(x.contiguous(), shard, state.enable_pdl) + return _EarlyFc(prologue, local, state.gather, current_capture_id()) + + +def _fused_side_chain( + worker, drafter, spec_metadata, num_gens: int, k: int, early: Optional[_EarlyFc] = None +) -> torch.Tensor: + """``ContextKvFork.side_chain`` with ``DFlashForCausalLM.precompute_context_kv`` written out. + + The same GEMMs (the K/V GEMM the drafter's own ``precompute_context_kv`` runs, + ``_kv_gemm``), norms, products and RoPE on the same operands in the same + order; the positions, columns and per-layer int32 RoPE positions come from + one prep kernel (``context_kv_prep``) instead of six torch kernels, and K and V + go into the arenas with one scatter (``context_kv_scatter``, V read in place + from the K/V GEMM's output) instead of ``v.contiguous()`` and two ``index_put``. + With ``early`` (``fire_early``), the prep kernel and the sharded ``fc`` GEMM already ran + after the last capture layer; this continues from the projection gather + (``hidden_norm(gather(local))``, the rest of ``sharded_project_target_hidden``). + """ + from tensorrt_llm._torch.models import modeling_dflash + + from ...ops.spec.context_kv_store import context_kv_scatter + + rows_per_req = k + 1 + rows = num_gens * rows_per_req + layers = drafter._num_attn_layers + nkv, hd = drafter._num_kv_heads, drafter._head_dim + if early is not None: + pro = early.prologue + projected = drafter.hidden_norm(early.gather.gather(early.local)) + else: + pro = _side_chain_prologue(worker, drafter, spec_metadata, num_gens, k) + projected = drafter.project_target_hidden(pro.gen_hs.reshape(-1, pro.gen_hs.shape[-1])) + slots, col, pos_i32, gen_hs = pro.slots, pro.col, pro.pos_i32, pro.gen_hs + # precompute_context_kv, its served branch (no input norm, the weight's dtype). + projected = projected.to(worker._ctx_k_buf.dtype) + kv_flat = _kv_gemm(drafter)(projected) + kv = kv_flat.view(rows, layers, 2, nkv, hd) + k_new = kv[:, :, 0].contiguous() + k_new = torch.nn.functional.rms_norm(k_new, (hd,), eps=drafter._k_norm_eps) + k_new = k_new * drafter._k_norm_stacked.view(1, layers, 1, hd) + # _fused_rope_inplace's flashinfer call, positions from the prep kernel. + k_flat = k_new.view(rows * layers, nkv * hd) + dummy_q = drafter._rope_dummy_q + if dummy_q is None or dummy_q.dtype != k_flat.dtype or dummy_q.shape[0] < k_flat.shape[0]: + dummy_q = k_flat.new_empty(k_flat.shape[0], drafter._head_dim) + drafter._rope_dummy_q = dummy_q + modeling_dflash._flashinfer_rope( + pos_i32, + dummy_q[: k_flat.shape[0]], + k_flat, + drafter._head_dim, + drafter._get_cos_sin_cache(), + drafter._is_neox, + ) + context_kv_scatter( + k_new, + kv_flat, + slots, + col, + worker._ctx_k_buf, + worker._ctx_v_buf, + rows_per_req, + worker._max_ctx, + ) + return gen_hs + + # -------------------------------------------------------------------------- # The worker # -------------------------------------------------------------------------- @@ -684,13 +1043,24 @@ def _reclass(spec_worker) -> None: spec_worker.__class__ = composed -def install(llm: nn.Module, core: nn.Module) -> None: +def install( + llm: nn.Module, + core: nn.Module, + *, + side_chain: bool = False, + early_fc: bool = False, + parallel_min_requests: Optional[int] = None, + declined: Optional[Callable[..., None]] = None, +) -> None: """Install the fork on the DFlash worker, the target core and the LM head (``POINT``). Idempotent; the shell calls it from every load hook after the other drafter features, so it composes with the worker class they left. A repeated install keeps the fork and its stream and takes the current - drafter. + drafter and lever settings: ``side_chain`` (``drafter_kv_side_chain``), + ``early_fc`` (``drafter_fc_early_fork``), ``parallel_min_requests`` + (``windowed_commit_parallel``'s MIN_REQUESTS, None for off) and ``declined`` + (``dispatch.Plan.declined``). Raises: ValueError: if the fork cannot be installed (the reason), so the @@ -712,6 +1082,10 @@ def install(llm: nn.Module, core: nn.Module) -> None: fork = ContextKvFork(worker, drafter, point) worker._kv_fork = fork fork.draft_model = drafter + fork.side_chain_on = bool(side_chain) + fork.early_fc_on = bool(side_chain and early_fc) + fork.parallel_min_requests = parallel_min_requests + fork.declined = declined core._context_kv_fork = fork if point == "head": lm_head = llm.lm_head diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py index 2e98c4ca23b6..97036f7cf0d9 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/drafter_block_decode/block_decode.py @@ -69,6 +69,56 @@ #: Operand dtypes the fused convolution takes; anything else runs the built-in. SUPPORTED_DTYPES = (torch.bfloat16, torch.float16, torch.float32) + +# ---- The drafter preamble computed by the prepare (drafter_preamble_from_commit) --------- +class DrafterPreamble: + """The integer tensors and the residual copy ``windowed_block_decode`` would derive itself. + + Published by the verify fold's prepare (``verify_fold/fold.py``, ``_fold_preamble``) + right before the drafter forward, taken by that forward when its inputs are these very + tensors (identity), otherwise ignored: ``i32s`` = (``num_ctx.to(int32)``, + that ``+ block``, ``slots.to(int32)``, ``query_positions.reshape(-1).to(int32)``), + ``residual`` a copy of ``noise``'s bytes (or None). + """ + + __slots__ = ("num_ctx", "slots", "query_positions", "noise", "i32s", "residual") + + def __init__(self, num_ctx, slots, query_positions, noise, i32s, residual): + self.num_ctx = num_ctx + self.slots = slots + self.query_positions = query_positions + self.noise = noise + self.i32s = i32s + self.residual = residual + + +_PUBLISHED_PREAMBLE = None + + +def publish_preamble(memo: DrafterPreamble) -> None: + global _PUBLISHED_PREAMBLE + _PUBLISHED_PREAMBLE = memo + + +def _take_preamble( + noise_embedding, query_positions, num_ctx_per_req, ctx_cache_batch_idx, B, block_size +): + """The published preamble when this forward's inputs are its tensors, else None; consumed either way.""" + global _PUBLISHED_PREAMBLE + memo, _PUBLISHED_PREAMBLE = _PUBLISHED_PREAMBLE, None + if ( + memo is None + or memo.noise is not noise_embedding + or memo.query_positions is not query_positions + or memo.num_ctx is not num_ctx_per_req + or memo.slots is not ctx_cache_batch_idx + or num_ctx_per_req.shape[0] != B + or query_positions.numel() != B * block_size + ): + return None + return memo + + #: Block rows on which the PDL chain's conv / merge entry triggers apply when the chain's #: install did not record their widths (``_pdl_entry_rows`` / ``_pdl_merge_rows``, #: ``chain.entry_rows`` / ``chain.MERGE_ENTRY_ROWS_BY_SM``): one request of K + 1 = 8 rows, whose @@ -311,10 +361,18 @@ def windowed_block_decode( block_size = noise_embedding.shape[1] hidden_states = noise_embedding - cache_seqlens_i32 = num_ctx_per_req[:B].to(torch.int32) - cache_batch_idx_i32 = ctx_cache_batch_idx.to(torch.int32) - seq_lens_after = cache_seqlens_i32 + block_size - query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) + preamble = _take_preamble( + noise_embedding, query_positions, num_ctx_per_req, ctx_cache_batch_idx, B, block_size + ) + if preamble is not None: + cache_seqlens_i32, seq_lens_after, cache_batch_idx_i32, query_positions_flat_i32 = ( + preamble.i32s + ) + else: + cache_seqlens_i32 = num_ctx_per_req[:B].to(torch.int32) + cache_batch_idx_i32 = ctx_cache_batch_idx.to(torch.int32) + seq_lens_after = cache_seqlens_i32 + block_size + query_positions_flat_i32 = query_positions.reshape(-1).to(torch.int32) residual = None has_block_conv = self.has_block_conv @@ -344,11 +402,15 @@ def windowed_block_decode( # The previous layer's hop already applied this layer's input norm. hs_normed_flat = prenormed elif residual is None: - residual = ( - hidden_states.clone(memory_format=torch.contiguous_format) - if hops is not None - else hidden_states.clone() - ) + if preamble is not None and preamble.residual is not None: + # The noise gather's second output, the same bytes as the clone. + residual = preamble.residual + else: + residual = ( + hidden_states.clone(memory_format=torch.contiguous_format) + if hops is not None + else hidden_states.clone() + ) hs_normed_flat = layer.input_layernorm(hs_flat) else: res_flat = residual.reshape(-1, residual.shape[-1]) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py index db66937ca604..936439e7b3a3 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/split_sampler/target_head.py @@ -100,6 +100,7 @@ from __future__ import annotations +import socket import time from typing import Optional @@ -196,6 +197,118 @@ def context_kv_head_fork(lm_head) -> tuple: return True, vars(lm_head).get("_context_kv_fork") +# ---- The logits all-gather over MULTIMEM (``head_multimem_gather``) ---- +#: With ``head_multimem_gather`` (and no ``lm_head_inbox``), ``TargetHead`` gathers the LM +#: head's shards of a ``MULTIMEM_GATHER_ROWS``-row call with +#: ``torch.ops.symm_mem.multimem_all_gather_out`` (one kernel, NVSwitch multicast stores, a +#: barrier before and after) into a symmetric-memory buffer instead of the NCCL ring, then +#: joins the shards with the built-in gather's own ``cat`` statement. A copy of the same bytes: +#: the logits are the built-in ones bit for bit. +MULTIMEM_GATHER_ROWS = GATHER_ROWS + + +def _torch_dist_world(mapping) -> None: + """A ``torch.distributed`` WORLD group over the MPI ranks, created once if none exists. + + Collective over the TP group (the host all-gather of rank 0's rendezvous address). Only + a group whose TP group is the whole MPI world is supported (``MultimemGather.reason``). + """ + import torch.distributed as dist + + from ... import dispatch as _dispatch + + if dist.is_initialized(): + return + # Rank 0's host and a port free on it right now (an inherited MASTER_PORT may be taken). + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("", 0)) + port = int(sock.getsockname()[1]) + try: + addr = socket.gethostbyname(socket.gethostname()) + except OSError: + addr = "127.0.0.1" + everyone = _dispatch.Distributed.get(mapping).tp_allgather((addr, port), small_payload=True) + addr, port = everyone[0] + dist.init_process_group( + backend="cuda:nccl,cpu:gloo", + init_method=f"tcp://{addr}:{port}", + rank=mapping.tp_rank, + world_size=mapping.tp_size, + device_id=torch.device("cuda", torch.cuda.current_device()), + ) + + +class MultimemGather: + """This rank's MULTIMEM gather of the LM head's ``[MULTIMEM_GATHER_ROWS, shard]`` BF16 shard.""" + + @staticmethod + def reason(mapping, lm_head) -> Optional[str]: + """Why the MULTIMEM gather cannot stand in for the head's all-gather here, or None (rank-local).""" + try: + import torch.distributed._symmetric_memory # noqa: F401 + except ImportError as e: + return f"no torch symmetric memory ({e})" + if not hasattr(torch.ops.symm_mem, "multimem_all_gather_out"): + return "torch has no symm_mem.multimem_all_gather_out" + from tensorrt_llm._utils import mpi_world_size + + if mpi_world_size() != mapping.tp_size or list(mapping.tp_group) != list( + range(mapping.tp_size) + ): + return "the TP group is not the whole MPI world" + if lm_head.lora is not None: + return "the LM head has LoRA" + if (int(lm_head.weight.shape[0]) * 2) % 4: + return "a shard row is not a whole number of 4-byte words" + return None + + def __init__(self, mapping, shard: int, device: torch.device) -> None: + """Collective over the TP group: the process group, the buffer and its rendezvous.""" + import torch.distributed as dist + import torch.distributed._symmetric_memory as symm_mem + + _torch_dist_world(mapping) + self.world = int(mapping.tp_size) + self.shard = int(shard) + self.group_name = str(dist.group.WORLD.group_name) + try: + symm_mem.enable_symm_mem_for_group(self.group_name) + except (AttributeError, RuntimeError): + pass # torch builds that enable it implicitly + self.buffer = symm_mem.empty( + (self.world * MULTIMEM_GATHER_ROWS, self.shard), dtype=torch.bfloat16, device=device + ) + handle = symm_mem.rendezvous(self.buffer, self.group_name) + self.multicast = bool(handle.multicast_ptr) + + def takes(self, hidden_states: torch.Tensor, lm_head) -> bool: + return ( + self.multicast + and hidden_states.shape[0] == MULTIMEM_GATHER_ROWS + and int(lm_head.weight.shape[0]) == self.shard + ) + + def __call__(self, lm_head: Linear, hidden_states: torch.Tensor) -> torch.Tensor: + """``Linear.forward`` of the column-parallel gathered head, with the MULTIMEM gather. + + ``apply_linear`` as the built-in forward calls it (so a fork point bound on the + instance fires), then ``_allgather``'s statements with the transport swapped: + ``view(-1, numel_base)``, the gather into ``[world * rows, shard]``, ``chunk`` and + ``cat`` along the last dim. + """ + from tensorrt_llm._torch.distributed.ops import get_output_info + + local = lm_head.apply_linear(hidden_states, lm_head.bias, None, None) + info = get_output_info(local, -1) + flat = local.contiguous().view(-1, info["numel_base"]) + # The NCCL op gathers ``flat`` along its dim 0: the ranks' blocks back to back. + out = self.buffer.view(-1)[: self.world * flat.numel()].view( + self.world * flat.shape[0], flat.shape[1] + ) + torch.ops.symm_mem.multimem_all_gather_out(flat, self.group_name, out) + return torch.cat([x.reshape(info["output_shape"]) for x in out.chunk(self.world)], dim=-1) + + class TargetHead: """The state the re-classed processor and worker share: the step's pending BF16 row.""" @@ -204,6 +317,14 @@ class TargetHead: #: This rank's ``[rows, shard]`` LM-head output while a vocab-parallel row is pending. pending_local: Optional[torch.Tensor] = None + #: ``head_multimem_gather``'s MULTIMEM gather, when every rank built it. + multimem: Optional["MultimemGather"] = None + #: Why ``head_multimem_gather`` asked for the gather and has none (agreed), or None. + multimem_refusal: Optional[str] = None + #: ``head_sampler_coop_fold``'s MAX_ROWS: head calls of at most this many rows run kernel A + #: with the soft cap computed by the whole CTA (0: never). + coop_max_rows: int = 0 + def __init__( self, processor: MuseSparkOpenLogitsProcessor, @@ -212,6 +333,8 @@ def __init__( mapping, gather: bool, vocab_parallel: bool = False, + multimem_gather: bool = False, + coop_max_rows: int = 0, ) -> None: self.multiplier = float(processor.logit_multiplier) self.soft_cap = float(processor.soft_cap) @@ -219,6 +342,7 @@ def __init__( self.gather = gather self.pdl = LAUNCH_PDL self.device = lm_head.weight.device + self.coop_max_rows = int(coop_max_rows) start = time.perf_counter() # Both loads are collective: every rank loads the library, then the group agrees, # so a rank whose load fails cannot leave the others waiting. @@ -242,6 +366,31 @@ def __init__( self.vp = _hs.VocabParallelWorkspace( mapping, int(lm_head.weight.shape[0]), self.max_rows ) + # head_multimem_gather: the MULTIMEM logits gather of MULTIMEM_GATHER_ROWS-row steps, + # when every rank can build it (collective). Its row band (8) is below + # vocab_parallel_head's (VP_MIN_ROWS..max_rows), so the two never take the same step + # (``produce``); the inbox gathers the same steps, so the plan never carries both. + self.multimem = None + self.multimem_refusal = None + if multimem_gather: + if gather: + raise ValueError( + "[musespark_open] head_multimem_gather: lm_head_inbox gathers the same steps" + ) + if self.vp is not None and MULTIMEM_GATHER_ROWS >= _hs.VP_MIN_ROWS: + raise ValueError( + "[musespark_open] head_multimem_gather: its row band overlaps " + "vocab_parallel_head's" + ) + refusal = _tp_agreement.agreed_refusal(mapping, MultimemGather.reason(mapping, lm_head)) + if refusal is None: + self.multimem = MultimemGather(mapping, int(lm_head.weight.shape[0]), self.device) + refusal = _tp_agreement.agreed_refusal( + mapping, None if self.multimem.multicast else "no multicast pointer" + ) + if refusal is not None: + self.multimem = None + self.multimem_refusal = refusal self.load_seconds = time.perf_counter() - start self.pending: Optional[torch.Tensor] = None self.pending_inbox = False @@ -348,6 +497,9 @@ def produce(self, hidden_states: torch.Tensor, lm_head: Linear) -> torch.Tensor: fork.fire("head") logits = self.workspace.gathered[:rows] self.pending_inbox = True + elif self.multimem is not None and self.multimem.takes(hidden_states, lm_head): + logits = self.multimem(lm_head, hidden_states) + self.pending_inbox = False else: logits = lm_head(hidden_states) self.pending_inbox = False @@ -434,6 +586,7 @@ def sample( self.soft_cap, self.states if logits.shape[0] <= _hs.MAX_ROWS else self.wide_states, self.pdl, + coop_max_rows=self.coop_max_rows, ) def vocab_parallel(self, logits) -> bool: @@ -738,12 +891,23 @@ def widen(llm) -> bool: return True -def install(llm, mapping, gather: bool, vocab_parallel: bool = False) -> None: - """Install the split sampler on the target's processor and DFlash worker. Idempotent. - - With ``gather`` (``lm_head_inbox``) the LM-head GEMM also writes into the multicast - inbox the sampler gathers; without it (``split_sampler``) the sampler takes the - built-in all-gather's logits. +def install( + llm, + mapping, + gather: bool, + vocab_parallel: bool = False, + multimem_gather: bool = False, + coop_max_rows: int = 0, +) -> "TargetHead": + """Install the split sampler on the target's processor and DFlash worker; returns the head. + + Idempotent. With ``gather`` (``lm_head_inbox``) the LM-head GEMM also writes into the + multicast inbox the sampler gathers; without it (``split_sampler``) the sampler takes + the built-in all-gather's logits, or, with ``multimem_gather`` + (``head_multimem_gather``), the 8-row steps' logits gathered over MULTIMEM (the head's + ``multimem``, or its ``multimem_refusal``). ``coop_max_rows`` + (``head_sampler_coop_fold``'s MAX_ROWS, 0 for off) is the largest head call whose + kernel A computes the soft cap with the whole CTA. Call after the drafter fast paths (``draft_head`` re-classes only the built-in worker class; ``fused_accept`` composes with a class derived from it). A fast path that @@ -757,7 +921,7 @@ def install(llm, mapping, gather: bool, vocab_parallel: bool = False) -> None: alike, before any other collective). """ if installed(llm): - return + return llm.logits_processor._split_head name = "lm_head_inbox" if gather else "split_sampler" reason = unsupported_reason(llm, mapping, gather) reason = _tp_agreement.agreed_refusal( @@ -766,7 +930,16 @@ def install(llm, mapping, gather: bool, vocab_parallel: bool = False) -> None: if reason is not None: raise ValueError(reason) processor, worker, lm_head = llm.logits_processor, llm.spec_worker, llm.lm_head - head = TargetHead(processor, lm_head, worker, mapping, gather, vocab_parallel) + head = TargetHead( + processor, + lm_head, + worker, + mapping, + gather, + vocab_parallel, + multimem_gather=multimem_gather, + coop_max_rows=coop_max_rows, + ) processor.__class__ = SplitSamplerLogitsProcessor processor._split_head = head composed = attach_worker(worker, head) @@ -778,6 +951,18 @@ def install(llm, mapping, gather: bool, vocab_parallel: bool = False) -> None: if gather else "" ) + + ( + f"; head_multimem_gather: the {MULTIMEM_GATHER_ROWS}-row steps' logits all-gathered " + "over MULTIMEM" + if head.multimem is not None + else "" + ) + + ( + f"; head_sampler_coop_fold: kernel A's soft cap by the whole CTA up to " + f"{head.coop_max_rows} rows" + if head.coop_max_rows + else "" + ) + ( f"; vocab_parallel_head: decode steps of {_hs.VP_MIN_ROWS}..{head.max_rows} rows sample " f"each rank's own shard (no logits all-gather), {head.vp.nbytes() / 2**20:.1f} MiB of " @@ -788,3 +973,4 @@ def install(llm, mapping, gather: bool, vocab_parallel: bool = False) -> None: + f"; library {head.library} (load {head.load_seconds:.2f} s); " + f"worker class {composed.__name__}" ) + return head diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py index bfbb099f3cde..bf65696e1d7f 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/kernels/verify_fold/fold.py @@ -69,7 +69,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Optional +from typing import Callable, Optional import torch @@ -79,16 +79,52 @@ from tensorrt_llm.logger import logger from ... import spec_worker as _drafter_worker -from ...ops.embedding.replicated_gather import replicated_embedding_gather +from ...ops.embedding.replicated_gather import ( + replicated_embedding_gather, + replicated_embedding_gather_twice, +) from ...ops.spec.dflash_accept import FORCE_FRACTIONAL from ...ops.spec.dflash_verify_fold import MAX_REQUESTS, VerifyFold, dflash_verify_fold from ...ops.spec.draft_path_walk_tail import draft_path_walk_tail from ..draft_head import head as _draft_head +from ..drafter_block_decode import block_decode as _block_decode from ..replicated_embedding import embedding as _replicated_embedding #: The drafter block width the fold writes positions for. MAX_BLOCK = 64 + +def _preamble_refusal(worker, fold, attn_metadata, num_gens: int, block_size: int) -> Optional[str]: + """Why the preamble commit cannot write this step's preamble, or None. + + With ``drafter_preamble_from_commit`` (``SpecLoopFoldDFlashWorker._fold_preamble``), on + a forked folded step the commit also writes the drafter forward's integer preamble + (``drafter_block_decode.DrafterPreamble``) and the ``_seq_lens_cuda`` fill, and the + noise gather also writes the residual's copy. Its kernel + (``context_kv_commit_preamble``) is the dense full-length arena's commit only: no KV + cache manager pages (``_ctx_paged`` / ``_ctx_block_tables``) and no windowed arena + (``_ctx_base``); every other arena, and every other shape, takes ``ContextKvFork.commit``. + """ + if worker._ctx_paged or worker._ctx_block_tables is not None or worker._ctx_base is not None: + return "the drafter context is not the dense full-length arena" + qpos = fold.query_positions + seq_lens = getattr(attn_metadata, "_seq_lens_cuda", None) + if not ( + qpos.dtype == torch.int64 + and qpos.is_contiguous() + and tuple(qpos.shape) == (num_gens, block_size) + ): + return "the fold's query positions are not [requests, block] contiguous int64" + if not ( + seq_lens is not None + and seq_lens.dtype in (torch.int32, torch.int64) + and seq_lens.is_contiguous() + and seq_lens.numel() >= num_gens + ): + return "the attention metadata's _seq_lens_cuda does not cover the requests" + return None + + #: Worker methods this module overrides; a worker class other than the #: built-in, ``draft_head``'s and ``fused_accept``'s that overrides one of them is refused. OVERRIDDEN = ( @@ -172,6 +208,10 @@ class SpecLoopFoldDFlashWorker(DFlashWorker): _fold_verify: bool = False _fold_path_walk_tail: bool = False + #: ``drafter_preamble_from_commit``: a forked folded step's commit writes the preamble. + _fold_preamble: bool = False + #: Reports a step a lever left to the path below it (``dispatch.Plan.declined``), or None. + _fold_declined: Optional[Callable[..., None]] = None #: Its ``_apply_dflash2_selector`` applies the draft head's deferred drafter #: logit transform (``draft_head.head._defers_transform``). _draft_transform_aware = True @@ -524,6 +564,48 @@ def _fold_forked_inputs( f"{num_gens} and K {spec_metadata.runtime_draft_len}" + ("; the step has no target features" if not has_target_features else "") ) + refusal = ( + _preamble_refusal(self, fold, attn_metadata, num_gens, block_size) + if self._fold_preamble + else None + ) + if self._fold_preamble and refusal is None: + # The commit also writes the drafter forward's integer preamble and the + # _seq_lens_cuda fill; the noise gather also writes the residual's copy. + num_ctx, i32s = fork.commit_with_preamble( + launched, + slots, + gen_num_accepted, + fold.query_positions, + attn_metadata._seq_lens_cuda, + total_tokens_per_req, + block_size, + ) + table = _replicated_embedding.table_of(embed_tokens) + residual = None + if table is not None and fold.noise_ids.numel() <= _replicated_embedding.MAX_ROWS: + noise_embed_2d, residual = replicated_embedding_gather_twice(table, fold.noise_ids) + else: + noise_embed_2d = self._fold_noise_block(embed_tokens, fold, num_gens, block_size) + fork.join(launched) + attn_metadata._seq_lens[0:num_gens] = total_tokens_per_req + inputs = self._drafter_inputs( + noise_embed_2d, fold.query_positions, num_ctx, slots, fold.gen_rows, fold.bonus + ) + # The dense arena (``_preamble_refusal``): the drafter reads the context by slot. + _block_decode.publish_preamble( + _block_decode.DrafterPreamble( + num_ctx, + inputs["ctx_cache_batch_idx"], + fold.query_positions, + noise_embed_2d, + i32s, + residual, + ) + ) + return inputs + if refusal is not None and self._fold_declined is not None: + self._fold_declined("drafter_preamble_from_commit", refusal, f"{num_gens} requests") num_ctx = fork.commit(launched, slots, gen_num_accepted) noise_embed_2d = self._fold_noise_block(embed_tokens, fold, num_gens, block_size) fork.join(launched) @@ -716,12 +798,21 @@ def unsupported_reason( return None -def install(spec_worker, verify_fold: bool, path_walk_tail: bool) -> bool: +def install( + spec_worker, + verify_fold: bool, + path_walk_tail: bool, + preamble: bool = False, + declined: Optional[Callable[..., None]] = None, +) -> bool: """Re-class the DFlash worker for the verify fold (``verify_fold``) and the path-walk tail (``path_walk_tail``) asked for; return whether anything changed. - Idempotent; the shell calls it from every drafter load hook, after - ``draft_head.install`` and ``drafter_worker.install``. + ``preamble`` (``drafter_preamble_from_commit``, with ``verify_fold``) has a forked + folded step's commit write the drafter preamble; ``declined`` reports the steps it + leaves to ``ContextKvFork.commit`` (``dispatch.Plan.declined``). Idempotent; the shell + calls it from every drafter load hook, after ``draft_head.install`` and + ``drafter_worker.install``. Raises: ValueError: if a fold asked for cannot be installed (``unsupported_reason``). @@ -738,13 +829,17 @@ def install(spec_worker, verify_fold: bool, path_walk_tail: bool) -> bool: reason = unsupported_reason(spec_worker, verify_fold, path_walk_tail) if reason is not None: raise ValueError(f"[musespark_open] {','.join(selected)} cannot run here: {reason}") + preamble = bool(preamble and verify_fold) if isinstance(spec_worker, SpecLoopFoldDFlashWorker): - changed = (spec_worker._fold_verify, spec_worker._fold_path_walk_tail) != ( - verify_fold, - path_walk_tail, - ) + changed = ( + spec_worker._fold_verify, + spec_worker._fold_path_walk_tail, + spec_worker._fold_preamble, + ) != (verify_fold, path_walk_tail, preamble) spec_worker._fold_verify = verify_fold spec_worker._fold_path_walk_tail = path_walk_tail + spec_worker._fold_preamble = preamble + spec_worker._fold_declined = declined return changed cls = type(spec_worker) composed = _COMPOSED_WORKERS.get(cls) @@ -761,6 +856,8 @@ def install(spec_worker, verify_fold: bool, path_walk_tail: bool) -> bool: spec_worker.__class__ = composed spec_worker._fold_verify = verify_fold spec_worker._fold_path_walk_tail = path_walk_tail + spec_worker._fold_preamble = preamble + spec_worker._fold_declined = declined sites = [ name for on, name in ( diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/embedding/replicated_gather.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/embedding/replicated_gather.py index 3cefd1131ae7..a5dc81f93706 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/embedding/replicated_gather.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/embedding/replicated_gather.py @@ -145,3 +145,53 @@ def is_valid(self, table: torch.Tensor, ids: torch.Tensor) -> None: replicated_embedding_gather = ReplicatedEmbeddingGather() + + +# ---- The gather into two outputs (drafter_preamble_from_commit's noise block) ------------- +@triton.jit +def _replicated_embedding_gather2_kernel( + table_ptr, + ids_ptr, + out_ptr, + copy_ptr, + vocab, + HIDDEN: tl.constexpr, + BLOCK: tl.constexpr, +): + # _replicated_embedding_gather_kernel's body, with the same bits also stored to copy_ptr. + row = tl.program_id(0) + chunk = tl.program_id(1) + token = tl.load(ids_ptr + row).to(tl.int64) + valid = (token >= 0) & (token < vocab) + cols = chunk * BLOCK + tl.arange(0, BLOCK) + in_row = cols < HIDDEN + source = tl.where(valid, token, 0) * HIDDEN + bits = tl.load(table_ptr + source + cols, mask=in_row & valid, other=0) + magnitude = bits.to(tl.int32) & 0x7FFF + bits = tl.where(magnitude == 0, tl.zeros_like(bits), bits) + tl.store(out_ptr + row.to(tl.int64) * HIDDEN + cols, bits, mask=in_row) + tl.store(copy_ptr + row.to(tl.int64) * HIDDEN + cols, bits, mask=in_row) + + +def replicated_embedding_gather_twice(table: torch.Tensor, ids: torch.Tensor): + """``(out, copy)``: ``replicated_embedding_gather(table, ids)`` and a second tensor with its bytes. + + One launch in place of the gather and a later ``out.clone()`` (the drafter's residual). + """ + vocab, hidden = table.shape + out = torch.empty((*ids.shape, hidden), dtype=table.dtype, device=table.device) + copy = torch.empty_like(out) + rows = ids.numel() + if rows == 0: + return out, copy + _replicated_embedding_gather2_kernel[(rows, triton.cdiv(hidden, BLOCK))]( + table.view(torch.int16), + ids.reshape(-1), + out.view(torch.int16), + copy.view(torch.int16), + vocab, + HIDDEN=hidden, + BLOCK=BLOCK, + num_warps=4, + ) + return out, copy diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu index b7aed7e11434..3694a5cb443a 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/_head_sampler/head_sampler.cu @@ -321,6 +321,108 @@ __global__ void __launch_bounds__(THREADS) splitFoldKernel(SplitFoldParams p) gridDependencyTrigger(); } +// ---- Cooperative kernel A (head_sampler_coop_fold) ---------------------------------------- +//! Kernel A on a plain BF16 row with the cap computed by every thread of the CTA. +//! +//! Grid (numRows, kSplitBlock / VTS); block THREADS >= VTS. The CTA owns virtual threads +//! vt0 .. vt0 + VTS - 1 (vt0 = blockIdx.y * VTS), i.e. the vectors v = vt0 + t + 1024 j. +//! Phase 1: all THREADS threads read those vectors (coalesced: VTS consecutive vectors per +//! j) and store each element's `__fmul_rn(capOf(logit), tempInv)` -- the value kernel A +//! hands to combineOnlineSoftmax -- in shared memory at [j][t]. Phase 2: thread t < VTS +//! folds its vectors j = 0, 1, ... in increasing j, each vector's 4 elements in order, +//! with the same combineOnlineSoftmax calls as splitFoldKernel. Every state is +//! splitFoldKernel's bit for bit: the same function of the same inputs in the same order; +//! only who computes the (pure, per-element) cap and where it waits changes. +//! +//! PDL: waits on its predecessor before its first global read; with EARLY_TRIGGER it +//! releases its dependent (kernel B, which itself waits for this grid's completion before +//! reading the states) right after that wait, so kernel B's launch overlaps this kernel. +template +__global__ void __launch_bounds__(THREADS) splitFoldCoopKernel(SplitFoldParams p) +{ + static_assert(THREADS >= VTS && THREADS % VTS == 0 && kSplitBlock % VTS == 0, "CTA shape"); + extern __shared__ float4 sCapped[]; + gridDependencyWait(); + if constexpr (EARLY_TRIGGER) + { + gridDependencyTrigger(); + } + int const row = blockIdx.x; + int const vt0 = blockIdx.y * VTS; + int const vecCount = p.fs.vocabSize / 4; + float const tempInv = loadRowParams(p.fs, row).tempInv; + uint2 const* plainRow = p.plain + static_cast(row) * vecCount; + // Vectors of the CTA's first virtual thread: the most any of its threads folds. + int const steps = vecCount > vt0 ? (vecCount - vt0 + kSplitBlock - 1) / kSplitBlock : 0; + for (int i = threadIdx.x; i < steps * VTS; i += THREADS) + { + int const j = i / VTS; + int const v = vt0 + (i - j * VTS) + j * kSplitBlock; + if (v < vecCount) + { + uint2 const packed = __ldg(plainRow + v); + __nv_bfloat16 const* elems = reinterpret_cast<__nv_bfloat16 const*>(&packed); + float4 x; + x.x = __fmul_rn(capOf(__bfloat162float(elems[0]), p.cap), tempInv); + x.y = __fmul_rn(capOf(__bfloat162float(elems[1]), p.cap), tempInv); + x.z = __fmul_rn(capOf(__bfloat162float(elems[2]), p.cap), tempInv); + x.w = __fmul_rn(capOf(__bfloat162float(elems[3]), p.cap), tempInv); + sCapped[i] = x; + } + } + __syncthreads(); + if (threadIdx.x < VTS) + { + // combineOnlineSoftmax(st, {x, 1, i}) in two passes over the thread's own slots. + // Pass 1 carries only the running max (and argmax): `keep = st.max >= x` picks the + // branch, and the branch's weight __expf(x - st.max) (keep) or __expf(st.max - x) + // depends on the max before the element only, so it is computed here, off the sum's + // chain, and stored over x with the branch in its sign bit (the weight is >= 0 or + // NaN; negation flips only the sign bit). Pass 2 runs the sum's chain alone: + // keep: st.sum + 1.0f * w; otherwise 1.0f + st.sum * w -- the two statements of + // combineOnlineSoftmax with b.sum = 1, in the same order. + int const vt = vt0 + threadIdx.x; + int const mine = vecCount > vt ? (vecCount - vt + kSplitBlock - 1) / kSplitBlock : 0; + float runMax = -FLT_MAX; + int runArg = 0; + for (int j = 0; j < mine; ++j) + { + int const base = (vt + j * kSplitBlock) * 4; + float4 x = sCapped[j * VTS + threadIdx.x]; + float* xs = reinterpret_cast(&x); +#pragma unroll + for (int e = 0; e < 4; ++e) + { + bool const keep = runMax >= xs[e]; + float const w = keep ? __expf(xs[e] - runMax) : __expf(runMax - xs[e]); + runArg = keep ? runArg : base + e; + runMax = keep ? runMax : xs[e]; + xs[e] = keep ? w : -w; + } + sCapped[j * VTS + threadIdx.x] = x; + } + float runSum = 0.0f; + for (int j = 0; j < mine; ++j) + { + float4 const x = sCapped[j * VTS + threadIdx.x]; + float const xs[4] = {x.x, x.y, x.z, x.w}; +#pragma unroll + for (int e = 0; e < 4; ++e) + { + bool const keep = (__float_as_uint(xs[e]) >> 31) == 0u; + float const w = fabsf(xs[e]); + runSum = keep ? runSum + 1.0f * w : 1.0f + runSum * w; + } + } + p.states[static_cast(row) * kSplitBlock + vt] = OnlineSoftmax{runMax, runSum, runArg}; + } + if constexpr (!EARLY_TRIGGER) + { + __syncthreads(); + gridDependencyTrigger(); + } +} + //! Kernel B: one VB-thread CTA per row (VB = splitVirtualBlock(numRows), the shipped //! block of the call). `params.logits` is the BF16 row (gathered or not); the soft cap is //! applied on every read. @@ -1277,6 +1379,88 @@ void foldPlain(at::Tensor const& logits, at::Tensor const& temperatures, at::Ten launchFold(p, threads, pdl); } +// ---- Cooperative kernel A launch (head_sampler_coop_fold) --------------------------------- +template +void launchFoldCoopShape(fs::SplitFoldParams const& p, bool pdl) +{ + auto kernel = fs::splitFoldCoopKernel; + int const vecCount = p.fs.vocabSize / 4; + size_t const smem = static_cast((vecCount + fs::kSplitBlock - 1) / fs::kSplitBlock) * VTS * sizeof(float4); + // Raised once per instantiation to the widest row this op admits (the walk's limit). + static bool const raised = [] + { + size_t const most = static_cast(fs::kSplitWalkVectors) * VTS * sizeof(float4); + C10_CUDA_CHECK(cudaFuncSetAttribute(fs::splitFoldCoopKernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(most))); + return true; + }(); + (void) raised; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(p.fs.numRows, fs::kSplitBlock / VTS); + config.blockDim = dim3(THREADS); + config.dynamicSmemBytes = smem; + config.stream = at::cuda::getCurrentCUDAStream(); + cudaLaunchAttribute attribute[1]; + attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attribute[0].val.programmaticStreamSerializationAllowed = pdl ? 1 : 0; + config.attrs = attribute; + config.numAttrs = 1; + C10_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel, p)); +} + +template +void launchFoldCoop(fs::SplitFoldParams const& p, int64_t vts, int64_t threads, bool pdl) +{ +#define FOLD_COOP_SHAPE(V, T) \ + if (vts == V && threads == T) \ + { \ + launchFoldCoopShape(p, pdl); \ + return; \ + } + FOLD_COOP_SHAPE(128, 1024) + FOLD_COOP_SHAPE(128, 512) + FOLD_COOP_SHAPE(64, 1024) + FOLD_COOP_SHAPE(64, 512) + FOLD_COOP_SHAPE(64, 256) + FOLD_COOP_SHAPE(32, 512) + FOLD_COOP_SHAPE(32, 256) +#undef FOLD_COOP_SHAPE + TORCH_CHECK( + false, "cooperative kernel A: unsupported (virtual threads per CTA, threads) = (", vts, ", ", threads, ")"); +} + +//! Cooperative kernel A on a plain BF16 row: the states of foldPlain, bit for bit. +void foldPlainCoop(at::Tensor const& logits, at::Tensor const& temperatures, at::Tensor const& states, + double multiplier, double invCap, double cap, int64_t vts, int64_t threads, bool earlyTrigger, bool pdl) +{ + TORCH_CHECK(logits.scalar_type() == at::kBFloat16, "logits must be BF16"); + at::cuda::CUDAGuard guard(logits.device()); + fs::SplitFoldParams p; + TORCH_CHECK(logits.dim() == 2 && logits.is_contiguous() && logits.size(1) % 4 == 0, "logits [rows, vocab]"); + TORCH_CHECK(reinterpret_cast(logits.data_ptr()) % 8 == 0, "logits must be 8-byte aligned"); + TORCH_CHECK(logits.size(1) / 4 <= fs::kSplitWalkVectors * fs::kSplitBlock, "vocabulary too wide"); + // The cooperative kernel writes the states of the kSplitBlock virtual block; calls above + // kWideBlockMaxRows rows take the narrow block (splitVirtualBlock) and run foldPlain. + TORCH_CHECK(fs::splitVirtualBlock(static_cast(logits.size(0))) == fs::kSplitBlock, + "cooperative kernel A takes at most ", fs::kWideBlockMaxRows, " rows"); + checkRowVector(temperatures, "temperatures", at::kFloat, logits.size(0)); + checkStates(states, logits.size(0)); + p.fs.temperatures = temperatures.data_ptr(); + p.fs.numRows = static_cast(logits.size(0)); + p.fs.vocabSize = static_cast(logits.size(1)); + p.plain = reinterpret_cast(logits.data_ptr()); + p.cap = softCap(multiplier, invCap, cap); + p.states = reinterpret_cast(states.data_ptr()); + if (earlyTrigger) + { + launchFoldCoop(p, vts, threads, pdl); + } + else + { + launchFoldCoop(p, vts, threads, pdl); + } +} + //! Kernel A with the gather: poll this rank's inbox copy, write `gathered`, reset the //! inbox; fold too when `fold`. void foldGather(at::Tensor const& inbox, at::Tensor const& gathered, at::Tensor const& temperatures, @@ -1602,6 +1786,9 @@ TORCH_LIBRARY(musespark_open_head_sampler, m) m.def( "draw(Tensor logits, Tensor temperatures, Tensor top_ks, Tensor top_ps, Tensor min_ps, Tensor? seed, " "Tensor? offset, Tensor states, float multiplier, float inv_cap, float cap, bool pdl) -> Tensor"); + m.def( + "fold_plain_coop(Tensor logits, Tensor temperatures, Tensor(a!) states, float multiplier, float inv_cap, " + "float cap, int vts, int threads, bool early_trigger, bool pdl) -> ()"); m.def("ref_fold(Tensor capped, Tensor temperatures, Tensor(a!) states) -> ()"); m.def("cap(Tensor x, float multiplier, float inv_cap, float cap) -> Tensor"); m.def("watch(Tensor buf, Tensor(a!) first_seen, int blocks, int timeout_ms) -> ()"); @@ -1629,6 +1816,7 @@ TORCH_LIBRARY_IMPL(musespark_open_head_sampler, CUDA, m) m.impl("fold_plain", &foldPlain); m.impl("fold_gather", &foldGather); m.impl("draw", &draw); + m.impl("fold_plain_coop", &foldPlainCoop); m.impl("ref_fold", &refFold); m.impl("cap", &capRows); m.impl("watch", &watch); diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit.py index f33f8e1db849..486e8be9e60e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_commit.py @@ -318,3 +318,163 @@ def is_valid( context_kv_commit = ContextKvCommit() + + +# ---- The commit with the drafter preamble's integer tensors (drafter_preamble_from_commit) +@triton.jit +def _context_kv_commit_preamble_kernel( + k_ptr, + v_ptr, + ctx_len_ptr, + slots_ptr, + accepted_ptr, + new_len_ptr, + num_ctx_ptr, + num_gens, + num_slots, + max_ctx, + dummy_slot, + k_slot_stride, + k_layer_stride, + k_col_stride, + v_slot_stride, + v_layer_stride, + v_col_stride, + num_ctx_i32_ptr, + seq_after_i32_ptr, + slots_i32_ptr, + qpos_ptr, + qpos_i32_ptr, + seq_lens_ptr, + seq_lens_value, + ROWS: tl.constexpr, + LAYERS: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, + PAIRS: tl.constexpr, + HAS_V: tl.constexpr, + DRAFT_BLOCK: tl.constexpr, + QBLOCK: tl.constexpr, +): + # _context_kv_commit_kernel's body on the dense arena (no base, last column max_ctx - 1), + # statement for statement. + for g in range(num_gens): + slot = tl.load(slots_ptr + g).to(tl.int64) + ctx = tl.load(ctx_len_ptr + slot) + accepted = tl.load(accepted_ptr + g).to(tl.int64) + for start in tl.static_range(0, ROWS * LAYERS, PAIRS): + k_cols, zeroed = _columns( + k_ptr, slot, ctx, accepted, max_ctx - 1, k_slot_stride, k_layer_stride, + k_col_stride, start, ROWS, LAYERS, WIDTH, BLOCK, PAIRS, + ) # fmt: skip + if HAS_V: + v_cols, _ = _columns( + v_ptr, slot, ctx, accepted, max_ctx - 1, v_slot_stride, v_layer_stride, + v_col_stride, start, ROWS, LAYERS, WIDTH, BLOCK, PAIRS, + ) # fmt: skip + k = tl.load(k_cols, mask=zeroed, other=0.0) + v = tl.load(v_cols, mask=zeroed, other=0.0) + tl.store(k_cols, _times_zero(k), mask=zeroed) + tl.store(v_cols, _times_zero(v), mask=zeroed) + else: + k = tl.load(k_cols, mask=zeroed, other=0.0) + tl.store(k_cols, _times_zero(k), mask=zeroed) + tl.store(new_len_ptr + g, ctx + accepted) + tl.debug_barrier() + for g in range(num_gens): + slot = tl.load(slots_ptr + g).to(tl.int64) + tl.store(ctx_len_ptr + slot, tl.load(new_len_ptr + g)) + for s in range(num_slots): + tl.store(ctx_len_ptr + s, tl.minimum(tl.load(ctx_len_ptr + s), max_ctx)) + # DFlashWorker._advance_ctx_len's _empty_dummy_slot: the padding / dummy rows' shared slot + # is emptied after the add, before the lengths the drafter reads. + tl.store(ctx_len_ptr + dummy_slot, tl.load(ctx_len_ptr + dummy_slot) * 0) + tl.debug_barrier() + for g in range(num_gens): + slot = tl.load(slots_ptr + g).to(tl.int64) + new_len = tl.load(ctx_len_ptr + slot) + tl.store(num_ctx_ptr + g, new_len) + # The drafter forward's preamble (windowed_block_decode): num_ctx_per_req.to(int32), + # + block_size, ctx_cache_batch_idx.to(int32); and the prepare's + # _seq_lens_cuda[0:gens] = K + 1. + new_len32 = new_len.to(tl.int32) + tl.store(num_ctx_i32_ptr + g, new_len32) + tl.store(seq_after_i32_ptr + g, new_len32 + DRAFT_BLOCK) + tl.store(slots_i32_ptr + g, slot.to(tl.int32)) + tl.store(seq_lens_ptr + g, seq_lens_value) + # query_positions.reshape(-1).to(int32), written by the verify fold before this launch. + q = tl.arange(0, QBLOCK) + qm = q < num_gens * DRAFT_BLOCK + tl.store(qpos_i32_ptr + q, tl.load(qpos_ptr + q, mask=qm).to(tl.int32), mask=qm) + + +def context_kv_commit_preamble( + k_arena, + v_arena, + ctx_len, + slots, + num_accepted, + num_ctx_out, + rows: int, + max_ctx: int, + query_positions, + seq_lens_cuda, + seq_lens_value: int, + draft_block: int, + i32s: tuple, + dummy_slot: int, +) -> None: + """The dense arena's ``context_kv_commit`` and ``DFlashWorker._commit_context``, writing the drafter preamble too. + + The zeroing of ``context_kv_commit`` (last column ``max_ctx - 1``, no base), then + ``_advance_ctx_len`` (the add, the clamp of every slot to ``max_ctx``, the emptied + ``dummy_slot``) and ``_drafter_context_lengths`` into ``num_ctx_out`` (``ctx_len[slots]``: + no block table on the dense arena). ``i32s = (num_ctx_i32 [gens], seq_lens_after_i32 [gens], slots_i32 [gens], + query_positions_i32 [gens * draft_block])`` (int32, allocated by the caller) get + the values ``windowed_block_decode`` derives from ``num_ctx_out``, ``slots`` and + ``query_positions`` (int64 ``[gens, draft_block]`` contiguous, written by the verify + fold before this launch); also sets ``seq_lens_cuda[:gens] = seq_lens_value``. + Integer copies and one add: the same values. + """ + num_gens = slots.numel() + device = ctx_len.device + layers, width = k_arena.shape[1], k_arena.shape[3] * k_arena.shape[4] + new_len = torch.empty(num_gens, dtype=torch.int64, device=device) + num_ctx_i32, seq_after, slots_i32, qpos_i32 = i32s + has_v = v_arena is not None + values = v_arena if has_v else k_arena + _context_kv_commit_preamble_kernel[(1,)]( + k_arena, + values, + ctx_len, + slots, + num_accepted, + new_len, + num_ctx_out, + num_gens, + ctx_len.numel(), + max_ctx, + dummy_slot, + k_arena.stride(0), + k_arena.stride(1), + k_arena.stride(2), + values.stride(0), + values.stride(1), + values.stride(2), + num_ctx_i32, + seq_after, + slots_i32, + query_positions, + qpos_i32, + seq_lens_cuda, + seq_lens_value, + ROWS=rows, + LAYERS=layers, + WIDTH=width, + BLOCK=triton.next_power_of_2(width), + PAIRS=tile_pairs(rows, layers, width), + HAS_V=has_v, + DRAFT_BLOCK=draft_block, + QBLOCK=triton.next_power_of_2(max(1, num_gens * draft_block)), + num_warps=4, + ) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_store.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_store.py new file mode 100644 index 000000000000..3c0241f3d699 --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/context_kv_store.py @@ -0,0 +1,292 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DFlash context-K/V side chain: its integer prologue and its arena store, one launch each. + +The ``context_kv_fork`` side chain (``kernels/context_kv_fork/fork.py``) stores a +decode step's ``K + 1`` captured rows of every generation request into the +drafter's dense context arena, unmasked. Around its GEMMs and norms the +built-in statements it reproduces launch about a dozen small torch kernels: + +* in front, ``ctx_len[slots]``, ``arange(K + 1)``, the positions ``ctx_len + r``, + their clamp to ``max_ctx - 1`` (the columns) and, inside + ``DFlashForCausalLM._fused_rope_inplace``, the positions as int32 repeated + per drafter layer for the RoPE kernel; +* behind the RoPE, ``v.contiguous()`` and the two ``index_put`` stores + ``k_arena[slot, :, col] = k`` / ``v_arena[slot, :, col] = v`` with their + value copies. + +``ContextKvPrep`` computes the positions, the columns and the repeated int32 +positions in one program; ``ContextKvScatter`` stores K (post norm and RoPE) +and V (read in place from the fused K/V GEMM's output) into both arenas in +one launch. Integer arithmetic and copies only: every value written is the +built-in's. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch._experimental.modeling_v2.catalog._op import Arch, Cell, OpWrapper + +#: Generation requests one prep call takes (``context_kv_commit.MAX_REQUESTS``). +MAX_REQUESTS = 64 + + +@triton.jit +def _ctx_kv_prep_kernel( + ctx_len_ptr, + slots_ptr, + positions_ptr, + col_ptr, + pos_i32_ptr, + num_gens, + max_ctx, + ROWS: tl.constexpr, + LAYERS: tl.constexpr, + BLOCK: tl.constexpr, +): + r = tl.arange(0, BLOCK) + rows = r < ROWS + for g in range(num_gens): + slot = tl.load(slots_ptr + g) + ctx = tl.load(ctx_len_ptr + slot) + pos = ctx + r.to(tl.int64) + tl.store(positions_ptr + g * ROWS + r, pos, mask=rows) + tl.store(col_ptr + g * ROWS + r, tl.minimum(pos, max_ctx - 1), mask=rows) + pos32 = pos.to(tl.int32) + for layer in tl.static_range(LAYERS): + tl.store(pos_i32_ptr + (g * ROWS + r) * LAYERS + layer, pos32, mask=rows) + + +class ContextKvPrep(OpWrapper): + """The side chain's positions, columns and per-layer int32 RoPE positions, in one program. + + ``ctx_len`` int64 ``[slots]``, ``slots`` int64 ``[gens]``; writes ``positions`` and + ``col`` int64 ``[gens * rows]`` and ``pos_i32`` int32 ``[gens * rows * layers]``: + + * ``positions[g * rows + r] = ctx_len[slots[g]] + r``; + * ``col = positions.clamp(max=max_ctx - 1)``; + * ``pos_i32 = positions.to(int32).repeat_interleave(layers)``. + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="the served decode step: one request, K + 1 = 8 rows, 5 drafter layers", + spec=dict(gens=1, rows=8, layers=5, slots=2), + ), + Cell( + why="a context at the arena's end: columns clamp, positions do not", + spec=dict(gens=1, rows=8, layers=5, slots=2, at_end=True), + ), + Cell( + why="several requests in distinct slots", + spec=dict(gens=3, rows=4, layers=2, slots=5), + ), + ) + + note = """ + Metadata consumed: none. One Triton program (grid 1, 1 warp) looping over + the requests. Call site: ``context_kv_fork.ContextKvFork.side_chain`` under + ``drafter_kv_side_chain``, first launch of the side stream. + Numerics: integer arithmetic; int32 is the low word, as ``.to(torch.int32)``. + Hard limits: rows <= 64, gens <= MAX_REQUESTS. + """ + + def __call__( + self, ctx_len, slots, positions, col, pos_i32, rows: int, layers: int, max_ctx: int + ): + _ctx_kv_prep_kernel[(1,)]( + ctx_len, + slots, + positions, + col, + pos_i32, + slots.numel(), + max_ctx, + ROWS=rows, + LAYERS=layers, + BLOCK=triton.next_power_of_2(rows), + num_warps=1, + ) + + def reference( + self, ctx_len, slots, positions, col, pos_i32, rows: int, layers: int, max_ctx: int + ): + offsets = torch.arange(rows, dtype=torch.long, device=ctx_len.device) + pos = ctx_len[slots].unsqueeze(1) + offsets.unsqueeze(0) + positions.copy_(pos.reshape(-1)) + col.copy_(pos.clamp(max=max_ctx - 1).reshape(-1)) + pos_i32.copy_(pos.reshape(-1).to(torch.int32).repeat_interleave(layers)) + + def is_valid( + self, ctx_len, slots, positions, col, pos_i32, rows: int, layers: int, max_ctx: int + ): + assert ctx_len.dtype == torch.int64 and slots.dtype == torch.int64 + n = slots.numel() * rows + assert ( + positions.dtype == torch.int64 and positions.numel() == n and positions.is_contiguous() + ) + assert col.dtype == torch.int64 and col.numel() == n and col.is_contiguous() + assert pos_i32.dtype == torch.int32 and pos_i32.numel() == n * layers + assert pos_i32.is_contiguous() + assert 1 <= rows <= 64 and slots.numel() <= MAX_REQUESTS + + +context_kv_prep = ContextKvPrep() + + +@triton.jit +def _ctx_kv_scatter_kernel( + k_ptr, + kv_ptr, + slots_ptr, + col_ptr, + k_arena_ptr, + v_arena_ptr, + max_ctx, + kv_row_stride, + k_slot_stride, + k_layer_stride, + k_col_stride, + v_slot_stride, + v_layer_stride, + v_col_stride, + ROWS: tl.constexpr, + LAYERS: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK_L: tl.constexpr, + BLOCK_W: tl.constexpr, + HAS_V: tl.constexpr, +): + n = tl.program_id(0) + g = n // ROWS + r = n % ROWS + slot = tl.load(slots_ptr + g) + col = tl.load(col_ptr + n) + # Rows whose columns clamp onto one column: the last of them is the one kept, as a + # deterministic index_put keeps the last duplicate (context_kv_commit's assumption). + if (r + 1 < ROWS) & (col >= max_ctx - 1): + return + layer = tl.arange(0, BLOCK_L) + w = tl.arange(0, BLOCK_W) + m = (layer < LAYERS)[:, None] & (w < WIDTH)[None, :] + layer64 = layer.to(tl.int64) + k = tl.load(k_ptr + n * (LAYERS * WIDTH) + layer[:, None] * WIDTH + w[None, :], mask=m) + if HAS_V: + # The fused K/V GEMM's row: per layer [K | V], each WIDTH wide. + v = tl.load( + kv_ptr + n * kv_row_stride + layer[:, None] * (2 * WIDTH) + WIDTH + w[None, :], mask=m + ) + k_dst = ( + k_arena_ptr + slot * k_slot_stride + layer64[:, None] * k_layer_stride + col * k_col_stride + ) + tl.store(k_dst + w[None, :], k, mask=m) + if HAS_V: + v_dst = ( + v_arena_ptr + + slot * v_slot_stride + + layer64[:, None] * v_layer_stride + + col * v_col_stride + ) + tl.store(v_dst + w[None, :], v, mask=m) + + +class ContextKvScatter(OpWrapper): + """Store a step's context K and V rows into the dense arenas, as the built-in ``index_put`` pair does. + + ``k`` BF16 ``[n, layers, kv_heads, head_dim]`` contiguous (post norm and RoPE); + ``kv`` the fused K/V GEMM output ``[n, layers * 2 * kv_heads * head_dim]`` whose + V half of layer ``l`` is the second ``kv_heads * head_dim`` block of the + layer's ``[K | V]`` pair; ``slots`` int64 ``[gens]``, ``col`` int64 + ``[gens * rows]`` (``n = gens * rows``). Writes, for ``i = g * rows + r``: + ``k_arena[slots[g], :, col[i]] = k[i]`` and + ``v_arena[slots[g], :, col[i]] = kv.view(n, layers, 2, kv_heads, head_dim)[i, :, 1]``. + Where rows of a request clamp onto the arena's last column, the last of them is + stored (a deterministic ``index_put``'s last-duplicate-wins). + """ + + ARCHS = frozenset({Arch.SM_103}) + + CELLS: tuple[Cell, ...] = ( + Cell( + why="the served decode step: one request, 8 rows, 5 layers, 1 kv head x 128", + spec=dict(gens=1, rows=8, layers=5, kv_heads=1, head_dim=128, slots=2), + ), + Cell( + why="a context at the arena's end: duplicate columns, the last row stored", + spec=dict(gens=1, rows=8, layers=5, kv_heads=1, head_dim=128, slots=2, at_end=True), + ), + Cell( + why="several requests, 2 kv heads x 64", + spec=dict(gens=3, rows=4, layers=2, kv_heads=2, head_dim=64, slots=5), + ), + ) + + note = """ + Metadata consumed: none. Grid ``gens * rows`` Triton programs (1 warp), each + storing one row's ``layers`` columns of both arenas. Call site: + ``context_kv_fork.ContextKvFork.side_chain`` under ``drafter_kv_side_chain``, + in place of ``v.contiguous()`` and the two ``index_put`` stores. + Numerics: copies. Hard limits: the (kv_heads, head_dim) elements of an arena + column contiguous; distinct slots across requests. + """ + + def __call__(self, k, kv, slots, col, k_arena, v_arena, rows: int, max_ctx: int): + n, layers = k.shape[0], k.shape[1] + width = k.shape[2] * k.shape[3] + has_v = v_arena is not None + values = v_arena if has_v else k_arena + _ctx_kv_scatter_kernel[(n,)]( + k, + kv, + slots, + col, + k_arena, + values, + max_ctx, + kv.stride(0), + k_arena.stride(0), + k_arena.stride(1), + k_arena.stride(2), + values.stride(0), + values.stride(1), + values.stride(2), + ROWS=rows, + LAYERS=layers, + WIDTH=width, + BLOCK_L=triton.next_power_of_2(layers), + BLOCK_W=triton.next_power_of_2(width), + HAS_V=has_v, + num_warps=1, + ) + + def reference(self, k, kv, slots, col, k_arena, v_arena, rows: int, max_ctx: int): + """The built-in statements (``fork.side_chain``), with the stores row by row in order.""" + n, layers, kv_heads, head_dim = k.shape + slot_flat = slots.unsqueeze(1).expand(-1, rows).reshape(-1) + v = kv.view(n, layers, 2, kv_heads, head_dim)[:, :, 1].contiguous() + for i in range(n): + k_arena[slot_flat[i], :, col[i]] = k[i] + if v_arena is not None: + v_arena[slot_flat[i], :, col[i]] = v[i] + + def is_valid(self, k, kv, slots, col, k_arena, v_arena, rows: int, max_ctx: int): + assert k.dim() == 4 and k.is_contiguous() and k.dtype == torch.bfloat16 + n, layers, kv_heads, head_dim = k.shape + assert kv.dim() == 2 and kv.shape == (n, layers * 2 * kv_heads * head_dim) + assert kv.stride(1) == 1 and kv.dtype == torch.bfloat16 + assert slots.dtype == torch.int64 and col.dtype == torch.int64 and col.numel() == n + assert n == slots.numel() * rows + for arena in (k_arena, v_arena): + if arena is None: + continue + assert arena.dim() == 5 and arena.dtype == torch.bfloat16 + assert arena.shape[1] == layers and arena.shape[3:] == (kv_heads, head_dim) + assert arena.stride(4) == 1 and arena.stride(3) == head_dim + + +context_kv_scatter = ContextKvScatter() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py index 08a01afdab61..78e2427b7db6 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/ops/spec/head_sampler.py @@ -101,6 +101,11 @@ INBOX_ROWS = 8 #: Threads per CTA of kernel A (8 CTAs per row). FOLD_THREADS = 128 +#: (virtual threads, threads) per CTA of the cooperative kernel A (``fold_plain_coop``, +#: early trigger) that ``SplitSample`` runs on calls of at most ``coop_max_rows`` rows +#: (``head_sampler_coop_fold``): every thread of the CTA computes the soft cap, and the +#: states are ``fold_plain``'s bit for bit on the 1024-thread virtual block. +COOP_FOLD_SHAPE = (32, 512) #: A poll waiting longer than this traps: a peer that never published. The trap ends the #: CUDA context (where the built-in NCCL all-gather would keep waiting), so any skew between #: two ranks' LM-head GEMMs above it -- a debugger or profiler pausing one rank, a host stall -- @@ -252,7 +257,10 @@ class SplitSample(OpWrapper): Metadata consumed: none. Two launches (kernel A grid (rows, V / 128) x 128 threads; kernel B grid rows x V threads, V = ``virtual_threads(rows)``: 1024 up to 128 rows, 512 above), both with programmatic dependent launch when - ``pdl``; each waits on its predecessor before its first global read. + ``pdl``; each waits on its predecessor before its first global read. A call + of at most ``coop_max_rows`` rows (and at most 128) runs kernel A as + ``fold_plain_coop`` (``COOP_FOLD_SHAPE``, early trigger), whose states are + ``fold_plain``'s bit for bit. Call site: the target DFlash worker's ``_sample_tokens_for_batch`` under the ``split_sampler`` fast path (``target_head.py``), on non-greedy batches with the fused @@ -293,11 +301,19 @@ def __call__( soft_cap: float, states: torch.Tensor, pdl: bool = False, + coop_max_rows: int = 0, ) -> torch.Tensor: m, inv, cap = cap_constants(multiplier, soft_cap) - torch.ops.musespark_open_head_sampler.fold_plain( - logits, temperatures, states, m, inv, cap, FOLD_THREADS, pdl - ) + if logits.shape[0] <= min(coop_max_rows, MAX_ROWS): + # Kernel A with the cap computed by the whole CTA: fold_plain's states on the + # 1024-thread virtual block (up to MAX_ROWS rows). + torch.ops.musespark_open_head_sampler.fold_plain_coop( + logits, temperatures, states, m, inv, cap, *COOP_FOLD_SHAPE, True, pdl + ) + else: + torch.ops.musespark_open_head_sampler.fold_plain( + logits, temperatures, states, m, inv, cap, FOLD_THREADS, pdl + ) return torch.ops.musespark_open_head_sampler.draw( logits, temperatures, top_ks, top_ps, min_ps, seed, offset, states, m, inv, cap, pdl ) @@ -315,10 +331,11 @@ def reference( soft_cap: float, states: torch.Tensor, pdl: bool = False, + coop_max_rows: int = 0, ) -> torch.Tensor: """The built-in head: the processor's soft cap, then the shipped fused sampler. - ``states`` and ``pdl`` are ignored. + ``states``, ``pdl`` and ``coop_max_rows`` are ignored. """ capped = reference_capped(logits, multiplier, soft_cap) return torch.ops.trtllm.fused_sample_from_logits( @@ -338,6 +355,7 @@ def is_valid( soft_cap: float, states: torch.Tensor, pdl: bool = False, + coop_max_rows: int = 0, ) -> None: if logits.dtype != torch.bfloat16 or logits.dim() != 2: raise ValueError("logits must be BF16 [rows, vocab]") @@ -348,6 +366,8 @@ def is_valid( raise ValueError("vocab must be a multiple of 4 and at most 4 * 64 * 1024") if states.numel() < state_words(rows): raise ValueError("states is smaller than new_states(rows)") + if coop_max_rows < 0: + raise ValueError("coop_max_rows must be 0 (off) or a row count") split_sample = SplitSample() diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/runtime.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/runtime.py index 49fb9e37a8d4..ddacac69b022 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/runtime.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/runtime.py @@ -30,6 +30,19 @@ store's four ``index_copy_`` and the three device-to-host snapshots become one launch that writes the store rows and the pinned host snapshots directly. +Two levers of the plan (``make_runtime_hooks`` reads ``model._plan``) change where and +how much of this runs: + +* ``step_loop_on_execution_stream`` declares + ``RuntimeHooks.step_loop_on_execution_stream``: the executor's overlap loop + then issues the sample launch and the next prepare launch on the execution + stream, right behind the decode graph, instead of on its thread's stream + behind two cross-stream waits per step. +* ``incremental_block_offsets`` makes the prepare launch copy + only the block-table columns whose host page index changed since its last + copy of that row (``_BlockOffsetsMemo``, which states when the memo is + trusted); the readable columns hold the same words as with the full copy. + Steps that do not qualify -- prefill or mixed batches, eager steps, graph capture, dummies, any configuration outside the one ``StepInputStaging.ineligible`` accepts -- run the built-in code untouched. Within a qualifying step, a write @@ -70,11 +83,114 @@ _INT_DTYPES = (torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8) +#: Changed-column runs per block-table row recorded as separate descriptors; a row with +#: more runs is re-copied as one span from its first to its last changed column. +_MAX_BLOCK_RUNS = 4 + +#: The reasons ``StepInputStaging.ineligible`` gives for every step of a configuration, not +#: for one step: under ``incremental_block_offsets`` each is reported as a decline of it. +CONFIG_DECLINES = frozenset( + { + "attention metadata type", + "draft KV cache manager", + "guided decoding", + "mrope", + "attention DP", + "overlap scheduler off", + "beam search", + "helix", + "pipeline parallelism", + } +) +#: The block-offset copy's own declines that hold for a configuration, reported likewise. +BLOCK_DECLINES = frozenset({"block offsets: layout or manager mode", "block offsets: cross KV"}) + def _addr(t: torch.Tensor, element_offset: int = 0) -> int: return t.data_ptr() + element_offset * t.element_size() +class _BlockOffsetsMemo: + """The host page indices the staged block-offset copies last wrote, per device row. + + A ``BLOCKS`` descriptor writes ``K[i] = f(page[i])`` and ``V[i] = f(page[i]) + + kv_offset`` into the device block table from the host page-index row, a pure + function of the page index, the pool's scale and its offset. So when a device + row ``(pool, seq)`` last received ``f`` of the host row ``row`` over columns + ``[0, n)`` with the same scale and offset, and nothing else has written the + device table since, re-copying a column whose page index is unchanged writes + the word already there; ``changed_runs`` lists the others. + + Ownership. The device table ``dst`` has two writers: this stager's launches + and the KV cache manager's own copy, which ``TrtllmAttentionMetadata.prepare`` + issues when no window records the write. The decoder runner opens a window + (``StepInputStaging.begin_if_eligible``) right before its only + ``_prepare_inputs`` call, so every native copy happens in a step the + stager saw: one it found ineligible, or a window that declined the copy or + was discarded. Each of those calls ``forget``. The memo is also reset when + the manager, the attention metadata object or the table's storage (pointer + and shape) is not the one of its last recording. The host row is read by + the launch from pinned memory after the host recorded it; the executor waits + for the step's prepare event before it changes the host page table again + (``_wait_for_decoder_input_copy``), so the values compared here are the + values the launch reads. + """ + + def __init__(self): + self._scope: Optional[tuple] = None + #: (pool, seq) -> (host row, scale, kv_offset, page indices of columns [0, n)). + self._rows: dict[tuple[int, int], tuple[int, int, int, np.ndarray]] = {} + + def forget(self) -> None: + """Drop everything: the next recording copies every readable column.""" + self._scope = None + self._rows.clear() + + def open(self, manager: Any, dst: torch.Tensor, attn_metadata: Any) -> None: + """Start a recording for ``dst``; resets the memo unless it last recorded this very table.""" + scope = (manager, attn_metadata, dst.data_ptr(), tuple(dst.shape)) + last = self._scope + if ( + last is None + or last[0] is not manager + or last[1] is not attn_metadata + or last[2:] != scope[2:] + ): + self._rows.clear() + self._scope = scope + + def changed_runs( + self, pool: int, seq: int, row: int, scale: int, kv_offset: int, pages: np.ndarray + ) -> list[tuple[int, int]]: + """Column runs ``[lo, hi)`` of ``pages`` (this step's first ``n`` page indices) to copy. + + Remembers ``pages`` as what the device row holds once the recorded copies have run. + """ + n = pages.shape[0] + entry = self._rows.get((pool, seq)) + if entry is None or entry[:3] != (row, scale, kv_offset): + runs = [(0, n)] if n else [] + else: + old = entry[3] + m = min(old.shape[0], n) + diff = np.flatnonzero(pages[:m] != old[:m]) + runs = [] + if diff.size: + cuts = np.flatnonzero(np.diff(diff) > 1) + starts = np.concatenate(([diff[0]], diff[cuts + 1])) + ends = np.concatenate((diff[cuts], [diff[-1]])) + 1 + runs = [(int(a), int(b)) for a, b in zip(starts, ends)] + if n > m: + if runs and runs[-1][1] == m: + runs[-1] = (runs[-1][0], n) + else: + runs.append((m, n)) + if len(runs) > _MAX_BLOCK_RUNS: + runs = [(runs[0][0], runs[-1][1])] + self._rows[(pool, seq)] = (row, scale, kv_offset, pages.copy()) + return runs + + class _PrepareStager: """One step's window, the model's ``StepInputStager``: records the prepare-phase device writes (see the module docstring). @@ -93,9 +209,15 @@ def __init__( runner: Any, attn_metadata: Any, spec_metadata: Any, + blocks_memo: Optional[_BlockOffsetsMemo] = None, + plan: Any = None, ): self.staging = staging self.counts = counts + #: When set, ``block_offsets`` re-copies only the changed columns (``_BlockOffsetsMemo``). + self.blocks_memo = blocks_memo + #: The plan the configuration-level block-copy declines are reported to, or None. + self.plan = plan self.runner = runner self.attn_metadata = attn_metadata self.declined: list[str] = [] @@ -323,6 +445,7 @@ def block_offsets( copies), which it is after an earlier decline. """ if self._after_decline(): + self._forget_blocks() return False host = getattr(manager, "host_kv_cache_block_offsets", None) if ( @@ -344,40 +467,86 @@ def block_offsets( or num_seqs > dst.shape[1] or beam_width > getattr(manager, "max_copy_beam_width", 0) ): - return self._decline("block offsets: layout or manager mode") + return self._decline_blocks("block offsets: layout or manager mode") from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp if getattr(manager, "kv_cache_type", None) == CacheTypeCpp.CROSS: - return self._decline("block offsets: cross KV") + return self._decline_blocks("block offsets: cross KV") tokens_per_block = int(manager.tokens_per_block) num_blocks = host.shape[3] cols = 0 for rid in request_ids: kv_cache = manager.kv_cache_map.get(rid) if kv_cache is None: - return self._decline("block offsets: request without a KV cache") + return self._decline_blocks("block offsets: request without a KV cache") cols = max(cols, math.ceil(int(kv_cache.capacity) / tokens_per_block) + 1) cols = min(cols, num_blocks) copy_idx = manager.index_mapper.get_copy_index(request_ids, num_contexts, beam_width, False) rows = copy_idx.tolist() if len(rows) != num_seqs: - return self._decline("block offsets: copy index length") - if not self.staging.room(host.shape[0] * len(rows)): - return self._decline("staging block full") + return self._decline_blocks("block offsets: copy index length") scales = [int(s) for s in manager._index_scale_ints] - kv_offsets = manager.kv_offset.tolist() + kv_offsets = [int(o) for o in manager.kv_offset.tolist()] host_rows, dst_rows = host.shape[1], dst.shape[1] + # (pool, seq, host row, column runs): one full-width run per row without a memo. + memo = self.blocks_memo + if memo is None: + plan = [ + (pool, seq, row, [(0, cols)]) + for pool in range(host.shape[0]) + for seq, row in enumerate(rows) + ] + else: + memo.open(manager, dst, self.attn_metadata) + pages = host.numpy() + plan = [ + ( + pool, + seq, + row, + memo.changed_runs( + pool, seq, row, scales[pool], kv_offsets[pool], pages[pool, row, 0, :cols] + ), + ) + for pool in range(host.shape[0]) + for seq, row in enumerate(rows) + ] + if not self.staging.room(sum(len(runs) for *_, runs in plan)): + return self._decline_blocks("staging block full") st = self.staging - for pool in range(host.shape[0]): - for seq, row in enumerate(rows): - src = _addr(host, (pool * host_rows + row) * 2 * num_blocks) - k = self._dst(_addr(dst, (pool * dst_rows + seq) * 2 * num_blocks), cols * 4) - v = self._dst(_addr(dst, ((pool * dst_rows + seq) * 2 + 1) * num_blocks), cols * 4) - st.blocks(k, v, src, cols, scales[pool], int(kv_offsets[pool])) + copied = 0 + for pool, seq, row, runs in plan: + src = _addr(host, (pool * host_rows + row) * 2 * num_blocks) + k = _addr(dst, (pool * dst_rows + seq) * 2 * num_blocks) + v = _addr(dst, ((pool * dst_rows + seq) * 2 + 1) * num_blocks) + for lo, hi in runs: + st.blocks( + self._dst(k + lo * 4, (hi - lo) * 4), + self._dst(v + lo * 4, (hi - lo) * 4), + src + lo * 4, + hi - lo, + scales[pool], + kv_offsets[pool], + ) + copied += hi - lo + if memo is not None: + self.counts["block columns copied"] += copied + self.counts["block columns kept"] += cols * len(plan) - copied self.block_cols = cols self.block_seqs = num_seqs return True + def _forget_blocks(self) -> None: + if self.blocks_memo is not None: + self.blocks_memo.forget() + + def _decline_blocks(self, reason: str) -> bool: + """Decline the block-offset copy: the caller's native copy rewrites the table, so the memo resets.""" + self._forget_blocks() + if self.blocks_memo is not None and self.plan is not None and reason in BLOCK_DECLINES: + self.plan.declined("incremental_block_offsets", reason) + return self._decline(reason) + def end(self, inputs: dict) -> None: """Close the window and launch what it recorded (``StepInputStager.end``). @@ -405,6 +574,7 @@ def discard(self) -> None: dropped (``StepInputStager.discard``).""" self._close() self._static = None + self._forget_blocks() self.counts["prepare aborted"] += 1 def static_inputs(self) -> Optional[tuple]: @@ -481,15 +651,25 @@ def _mirror_static_inputs(self, inputs: dict) -> Optional[tuple]: class StepInputStaging: """The model's ``StepInputStagerSource``: one ``_PrepareStager`` window per eligible step.""" - def __init__(self, pp_size: int = 1): + def __init__(self, pp_size: int = 1, incremental_blocks: bool = False, plan: Any = None): """Allocate the pinned staging block. Args: pp_size: the model's pipeline-parallel size; any value above 1 is ineligible. + incremental_blocks: re-copy only the changed block-table columns + (``_BlockOffsetsMemo``, the ``incremental_block_offsets`` lever); otherwise + every readable column, every step. + plan: the model's plan (``dispatch.Plan``): ``begin_if_eligible`` counts its + steps (``Plan.step``) and reports the steps ``incremental_block_offsets`` + declined (``Plan.declined``); None reports nothing. """ self.pp_size = pp_size + self.plan = plan self.staging = _sc.StepStaging() self.counts: Counter = Counter() + self.blocks_memo: Optional[_BlockOffsetsMemo] = ( + _BlockOffsetsMemo() if incremental_blocks else None + ) def ineligible( self, @@ -584,7 +764,14 @@ def begin_if_eligible( enable_spec_decode=enable_spec_decode, promoted=promoted, ) + if self.plan is not None: + self.plan.step() if reason is not None: + # The native prepare of this step rewrites the block table. + if self.blocks_memo is not None: + self.blocks_memo.forget() + if reason in CONFIG_DECLINES and self.plan is not None: + self.plan.declined("incremental_block_offsets", f"step input staging: {reason}") key = f"prepare declined: {reason}" if not self.counts[key]: # Once per reason, so that a configuration change that declines every @@ -600,7 +787,15 @@ def begin_if_eligible( self.counts[key] += 1 return None self.staging.reset() - return _PrepareStager(self.staging, self.counts, runner, attn_metadata, spec_metadata) + return _PrepareStager( + self.staging, + self.counts, + runner, + attn_metadata, + spec_metadata, + self.blocks_memo, + self.plan, + ) class SampleStoreWriter: @@ -872,20 +1067,30 @@ def make_runtime_hooks(model: Any) -> RuntimeHooks: exists before any CUDA graph is captured. Args: - model: the shell; both parts get its ``_mapping.pp_size`` (1 without one). + model: the shell; both parts get its ``_mapping.pp_size`` (1 without one). Its plan + (``_plan``) selects the parts (``deferred_input_prep``, ``sampler_store``) and the + levers ``step_loop_on_execution_stream`` and ``incremental_block_offsets``; a + model without one gets both parts and neither lever. Returns: - ``RuntimeHooks`` declaring both the step-input stager - (``StepInputStaging``) and the sampler store writer (``SampleStoreWriter``). + ``RuntimeHooks`` declaring the step-input stager (``StepInputStaging``) and the + sampler store writer (``SampleStoreWriter``) the plan carries. """ pp_size = int(getattr(getattr(model, "_mapping", None), "pp_size", 1) or 1) + plan = getattr(model, "_plan", None) + stager = plan is None or plan.enabled("deferred_input_prep") + writer = plan is None or plan.enabled("sampler_store") + one_stream = plan is not None and plan.enabled("step_loop_on_execution_stream") + incremental_blocks = plan is not None and plan.enabled("incremental_block_offsets") hooks = RuntimeHooks( - step_input_stager=StepInputStaging(pp_size), - sampler_store_writer=SampleStoreWriter(pp_size), + step_input_stager=StepInputStaging(pp_size, incremental_blocks, plan) if stager else None, + sampler_store_writer=SampleStoreWriter(pp_size) if writer else None, + step_loop_on_execution_stream=one_stream, ) _sc.warm_up() logger.info( - f"[musespark_open] executor hooks: step input staging and sampler store writer, " - f"pp_size={pp_size}" + f"[musespark_open] executor hooks: step input staging {'on' if stager else 'off'}, " + f"sampler store writer {'on' if writer else 'off'}, pp_size={pp_size}, step loop on " + f"the execution stream={one_stream}, incremental block offsets={incremental_blocks}" ) return hooks From 076269f078e3c55fdc2bf175840f8244360011ba Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:11:06 +0000 Subject: [PATCH 10/15] [None][feat] modeling_v2 musespark_open tep: a lever resolver that picks the fast paths per deployment New levers.py: one table of levers with a rule each. dispatch.compose resolves it once per process from the routing key and the deployment it sees: device, tensor-parallel width, max_batch_size, max_num_tokens, the serving role (aggregated, or the context / generation side of a disaggregated server) and the drafter's context arena (dense, windowed or paged). Installers read their decision and thresholds from the plan; no other module of the target reads the environment. - Lever kinds: composition and bitwise levers follow their rule; numeric levers also need their deployment class; lossy levers are never on unless requested; flagged levers stay off until requested. - Knobs: MUSESPARK_OPEN_LEVER_=auto|on|off per lever, declared parameters MUSESPARK_OPEN_LEVER___, MUSESPARK_OPEN_LEVER__ROLE and MUSESPARK_OPEN_LEVER__DEFAULT=auto|baseline. Unset means auto. baseline selects the levers the target served before this change, as a control. Unknown names, bad values, out-of-range parameters, a request a structural rule excludes and disagreement across ranks raise before anything is installed. - Structural rules: the drafter arena selects the commit and side-chain paths; TRTLLM_KV_DEFER_CLOSE=1 excludes the execution-stream loop; UCX_RNDV_THRESH may not be set on a disaggregated role. - The resolved and the installed lever tables are logged once on rank 0, the installed one as JSON after the last install hook, and per-step declines are counted. - The modeling_v2 router records the engine's deployment facts per target class. Host-only tests cover the grammar, the overrides, the structural rules, baseline reproduction and the per-deployment lever tables. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../modeling_v2/_router_index.py | 47 +- .../dispatch.py | 582 ++++++-- .../drafter.py | 29 +- .../layers.py | 256 +++- .../levers.py | 1138 ++++++++++++++ .../modeling.py | 200 ++- .../weights.py | 45 +- .../test_modeling_v2_drafter_context.py | 26 +- .../test_modeling_v2_levers.py | 1318 +++++++++++++++++ 9 files changed, 3418 insertions(+), 223 deletions(-) create mode 100644 tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py create mode 100644 tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py index 1fe1bb1135bb..409d6d4e4eac 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py @@ -26,7 +26,8 @@ are deliberately not routing inputs. They do not change which target runs, so a target claims a configuration regardless of them; ``require`` promises that the target ran, not that a particular knob setting was measured. The resolver records -``max_batch_size`` for the target it returns (``resolved_max_batch_size``), which a +``max_batch_size`` for the target it returns (``resolved_max_batch_size``), and the +engine facts a target's own rules read (``resolved_facts``: ``DeploymentFacts``), which a target may read to size or leave out a fast path inside itself. Routing modules are imported lazily: resolving a GptOss config never imports @@ -36,6 +37,7 @@ from __future__ import annotations +import dataclasses import enum import os from importlib import import_module @@ -137,6 +139,48 @@ def resolved_max_batch_size(target: type) -> Optional[int]: return _RESOLVED_MAX_BATCH_SIZES.get(target) +@dataclasses.dataclass(frozen=True) +class DeploymentFacts: + """The engine settings ``modeling_v2_resolve`` saw for a target, for the target's own rules. + + Not routing inputs. ``disaggregated`` is whether the engine carries a cache transceiver + (a context or generation server of a disaggregated deployment); ``graph_batch_sizes`` + the CUDA graph batch sizes it captures (None when it captures none or lists none). + """ + + max_batch_size: Optional[int] + max_num_tokens: Optional[int] + disaggregated: bool + graph_batch_sizes: Optional[Tuple[int, ...]] + + +# The engine facts each target class was last resolved under. +_RESOLVED_FACTS: Dict[type, DeploymentFacts] = {} + + +def resolved_facts(target: type) -> Optional[DeploymentFacts]: + """The ``DeploymentFacts`` ``modeling_v2_resolve`` saw for ``target``; None for a class built without them.""" + return _RESOLVED_FACTS.get(target) + + +def _facts_of(llm_args: "TorchLlmArgs") -> DeploymentFacts: + def count(name: str) -> Optional[int]: + value = getattr(llm_args, name, None) + return None if value is None else int(value) + + graphs = getattr(llm_args, "cuda_graph_config", None) + sizes = getattr(graphs, "batch_sizes", None) if graphs is not None else None + # The executor builds a cache transceiver only for a config with a backend + # (``disaggregation.kv_cache_transceiver``). + transceiver = getattr(llm_args, "cache_transceiver_config", None) + return DeploymentFacts( + max_batch_size=count("max_batch_size"), + max_num_tokens=count("max_num_tokens"), + disaggregated=transceiver is not None and getattr(transceiver, "backend", None) is not None, + graph_batch_sizes=tuple(int(size) for size in sizes) if sizes else None, + ) + + def routing_module(arch: str): """Import and return the routing module for ``arch``, or None.""" name = MODELING_V2_ROUTERS.get(arch) @@ -196,4 +240,5 @@ def modeling_v2_resolve(config: "ModelConfig", llm_args: "TorchLlmArgs") -> Opti _RESOLVED_MAX_BATCH_SIZES.pop(target, None) else: _RESOLVED_MAX_BATCH_SIZES[target] = int(max_batch_size) + _RESOLVED_FACTS[target] = _facts_of(llm_args) return target diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py index 206141714615..34a466420572 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/dispatch.py @@ -2,22 +2,28 @@ # SPDX-License-Identifier: Apache-2.0 """Which fast paths a model of this target is built with, decided once at construction. -``FAST_PATHS`` names every fast path of the target and the fast paths each one -builds on (``depends_on``). A composition (``TargetPaths``) is the set of fast -paths served on one GPU architecture at one tensor-parallel size (its -``CompositionKey``), closed under ``depends_on``, with the GPU topology it was -validated on; it applies when the drafter is DFlash drafting a linear chain of -``DRAFT_LEN`` tokens on an NVFP4 checkpoint. Those facts are structural: a -deployment that misses one has no composition, and runs the built-in forward -under ``modeling_v2="auto"`` (``NO_PATHS``) or raises under ``require``. Nothing -reads the environment to pick or trim a composition. - -``compose`` then checks the public preconditions the planned fast paths are -written against (``PRECONDITIONS``, and the composition's own), once, agreed -over the TP group. A missing one, on any rank, raises ``ValueError`` naming each -one, under ``auto`` as under ``require``: a composition is never served in part -and never dropped because of a setting. The model's installers ask the -resulting ``Plan`` per fast path (``Plan.enabled``). +``FAST_PATHS`` names every lever of the target -- each fast path, and the refinements +and installs registered beside them -- and the levers each one builds on +(``depends_on``). A composition (``TargetPaths``) is the set of levers served on one +GPU architecture at one tensor-parallel size (its ``CompositionKey``), closed under +``depends_on``, with the GPU topology it was validated on; it applies when the +drafter is DFlash drafting a linear chain of ``DRAFT_LEN`` tokens on an NVFP4 +checkpoint. Those facts are structural: a deployment that misses one has no +composition, and runs the built-in forward under ``modeling_v2="auto"`` +(``NO_PATHS``) or raises under ``require``. + +``compose`` resolves the composition's lever table (``levers.resolve``): the +structural rules of the deployment's drafter arena and facts +(``structural_rules``), then each lever's knob and deployment rule. The lever +table is the only environment ``compose`` reads (``levers.environment``). It then +checks the public preconditions the planned fast paths are written against +(``PRECONDITIONS``, and the composition's own), once, agreed over the TP group +together with the arena and the lever environment. A missing precondition, a lever +error or a disagreement, on any rank, raises ``ValueError`` on every rank, under +``auto`` as under ``require``: a composition is never served in part because of a +missing setting. The model's installers ask the resulting ``Plan`` per lever +(``Plan.enabled``, ``Plan.param``) and report what they installed and what a step +declined (``Plan.installed``, ``Plan.declined``). """ from __future__ import annotations @@ -44,8 +50,10 @@ from tensorrt_llm._torch.distributed.communicator import Distributed from tensorrt_llm.logger import logger +from . import levers as _levers + #: The draft tokens per step of the one speculative-decoding mode the compositions serve. -DRAFT_LEN = 7 +DRAFT_LEN = _levers.DRAFT_LEN #: The log tag of this family. TAG = "[musespark_open]" @@ -57,7 +65,7 @@ @dataclasses.dataclass(frozen=True) class FastPath: - """One fast path: what it replaces, the fast paths it builds on, and when it installs. + """One lever of this target: what it changes, the levers it builds on, and when it installs. ``hook`` names the shell hook that installs it: ``construct`` (the shell's constructor, on the built model), ``load`` (after the target weights load), @@ -249,6 +257,12 @@ class FastPath: "expert finalize all-reduce on the fused hop's transport", "construct", ), + FastPath( + "wide_finalize_transport", + "expert finalize all-reduce of 17..32-token calls on a 32-token TP4 transport layout", + "construct", + ("finalize_transport",), + ), FastPath( "router_gemm", "router logits of 9..512-row calls as a BF16x3 split-K tensor-core GEMM", @@ -413,13 +427,140 @@ class FastPath: # Executor. FastPath("deferred_input_prep", "a decode step's input writes as one launch", "executor"), FastPath("sampler_store", "the sampler's store update and snapshots as one launch", "executor"), + # Levers on the paths above (``LEVER_NAMES``): each changes how the path it builds on + # runs, or is an install the plan carries; their rules are ``levers.RULES``. + FastPath( + "packed_prep_append", + "multi-request decode prep-append calls of MIN_ROWS rows or more on the packed kernel", + "construct", + ("multi_request_prep_append",), + ), + FastPath( + "windowed_commit_parallel", + "windowed context-K/V commits of MIN_REQUESTS requests or more as a parallel zeroing " + "and a tuned compaction", + "draft", + ("context_kv_fork",), + ), + FastPath( + "context_fp8_dense", + "the dense projections of context calls of MIN_ROWS rows or more on the Rubin MXFP8 " + "GEMM (lossy)", + "load", + ), + FastPath( + "context_fp8_attn_gather", + "the sequence-parallel attention input all-gathered in MXFP8 (lossy)", + "load", + ("context_fp8_dense", "context_sequence_parallel"), + ), + FastPath( + "wide_direct_transports", + "the pool and unpool direct transports of 16..32-row captured decode steps", + "construct", + ("wide_prologue", "pool_into_route", "unpool_into_boundary"), + ), + FastPath( + "wide_router_per_row", + "the router shard of 16..32-row decode steps with each row's one-request logits", + "construct", + ("wide_prologue",), + ), + FastPath( + "wide_oproj_publish", + "o_proj of 16..32-row decode steps publishing into the hop's 32-row direct transport", + "construct", + ("attention_output_into_hop", "wide_post_attention_hop"), + ), + FastPath( + "graph_start_qkv_prefetch", + "layer 0's QKV weight prefetch forked at the start of a decode forward", + "construct", + ("next_qkv_prefetch",), + ), + FastPath( + "step_loop_on_execution_stream", + "the executor's overlap loop issuing the sample and prepare launches on the execution " + "stream", + "executor", + ("deferred_input_prep",), + ), + FastPath( + "incremental_block_offsets", + "the staged block-offset copy of only the columns whose page index changed", + "executor", + ("deferred_input_prep",), + ), + FastPath( + "drafter_kv_side_chain", + "the context-K/V side chain with its prologue and arena store fused (dense arena)", + "draft", + ("context_kv_fork",), + ), + FastPath( + "drafter_fc_early_fork", + "the side chain's local fc GEMM forked after the last DFlash capture layer", + "draft", + ("drafter_kv_side_chain", "drafter_fc_shard"), + ), + FastPath( + "drafter_preamble_from_commit", + "the drafter forward's integer preamble written by the folded step's commit (dense arena)", + "draft", + ("verify_fold", "context_kv_fork"), + ), + FastPath( + "head_sampler_coop_fold", + "the head sampler's kernel A with the soft cap computed by the whole CTA, head calls " + "of MAX_ROWS rows or fewer", + "draft", + ("split_sampler",), + ), + FastPath( + "head_multimem_gather", + "the 8-row head logits all-gathered over NVLink multicast (MULTIMEM)", + "draft", + ("split_sampler",), + ), + FastPath( + "prefill_cutedsl_experts", + "the NVFP4 experts of 4096..8192-row calls on the Rubin CuTe DSL GEMMs", + "load", + ), ) #: The name of every fast path. FAST_PATH_NAMES: FrozenSet[str] = frozenset(path.name for path in FAST_PATHS) +#: The levers registered beside the fast paths (the last group of ``FAST_PATHS``). +LEVER_NAMES: FrozenSet[str] = frozenset( + { + "packed_prep_append", + "windowed_commit_parallel", + "context_fp8_dense", + "context_fp8_attn_gather", + "wide_direct_transports", + "wide_router_per_row", + "wide_oproj_publish", + "graph_start_qkv_prefetch", + "step_loop_on_execution_stream", + "incremental_block_offsets", + "drafter_kv_side_chain", + "drafter_fc_early_fork", + "drafter_preamble_from_commit", + "head_sampler_coop_fold", + "head_multimem_gather", + "prefill_cutedsl_experts", + } +) +#: The fast paths a composition is built from; a composition lists its levers explicitly. +BASE_PATH_NAMES: FrozenSet[str] = FAST_PATH_NAMES - LEVER_NAMES + _BY_NAME: Dict[str, FastPath] = {path.name: path for path in FAST_PATHS} +#: Every lever and the levers it builds on, in ``FAST_PATHS`` order (``levers.resolve``). +DEPENDS_ON: Mapping[str, Tuple[str, ...]] = {path.name: path.depends_on for path in FAST_PATHS} + def closure(names: Iterable[str]) -> FrozenSet[str]: """``names`` and every fast path they build on, transitively. @@ -519,11 +660,17 @@ def __post_init__(self) -> None: raise ValueError(f"{TAG} a served fast path builds on unserved {missing}") +#: Steps between two log lines of the lever decline counts (``Plan.step``). +DECLINE_REPORT_STEPS = 10_000 + + @dataclasses.dataclass(frozen=True) class Plan: - """The fast paths a model is built with: a whole composition, or none (``NO_PATHS``). + """The levers a model is built with: a whole composition's table, or none (``NO_PATHS``). - Closed under ``depends_on`` like a composition (``plan_of`` closes a set). + Closed under ``depends_on`` like a composition (``plan_of`` closes a set). ``levers`` + is the table ``compose`` resolved (None for ``NO_PATHS`` and for a plan built from + names); ``rank0`` says whether this rank logs the per-server lines. Raises: ValueError: at construction, for a name that is not a fast path or a fast @@ -531,6 +678,18 @@ class Plan: """ paths: FrozenSet[str] = frozenset() + levers: Optional[_levers.LeverTable] = dataclasses.field(default=None, compare=False) + rank0: bool = dataclasses.field(default=True, compare=False) + #: Lever -> (installed, why not), as the installers reported (``installed``). + reports: Dict[str, Tuple[bool, str]] = dataclasses.field( + default_factory=dict, compare=False, repr=False + ) + #: (lever, reason) -> steps declined (``declined``). + declines: Dict[Tuple[str, str], int] = dataclasses.field( + default_factory=dict, compare=False, repr=False + ) + #: Steps counted by ``step``, in a one-element list so the frozen plan can count. + clock: List[int] = dataclasses.field(default_factory=lambda: [0], compare=False, repr=False) def __post_init__(self) -> None: unknown = sorted(set(self.paths) - FAST_PATH_NAMES) @@ -550,6 +709,65 @@ def enabled(self, path: str) -> bool: raise KeyError(f"{TAG} not a fast path: {path!r}") return path in self.paths + def param(self, name: str, key: str) -> Any: + """Lever ``name``'s parameter ``key``: as resolved, else its declared default. + + Raises: + KeyError: for a name that is not a fast path, or a parameter ``name`` does + not declare. + """ + if name not in FAST_PATH_NAMES: + raise KeyError(f"{TAG} not a fast path: {name!r}") + if self.levers is not None: + return self.levers.param(name, key) + return _levers.default_param(name, key) + + def installed(self, name: str, ok: bool, why: str = "") -> None: + """An installer's report on lever ``name``: installed, or refused for ``why``. + + The installer passes an outcome agreed over the TP group, so every rank reports + the same. A refusal of a lever its knob set ``on`` raises; any other refusal is + a warning, and the installed table (``report``) shows it. + + Raises: + ValueError: for a refusal of a lever set ``on`` by its knob. + """ + self.reports[name] = (bool(ok), why) + if ok: + return + decision = self.levers.decisions.get(name) if self.levers is not None else None + message = f"{TAG} lever {name} is not installed: {why}" + if decision is not None and decision.source == "knob": + raise ValueError(f"{message} ({_levers.knob_name(name)}=on)") + if self.rank0: + logger.warning(f"{message}; the model runs without it") + + def declined(self, name: str, reason: str, shape: str = "") -> None: + """A step lever ``name`` left to the path below it, for ``reason``. + + Counted per (lever, reason); the first of each is logged as a warning with the + step's ``shape``, and ``step`` logs the counts. + """ + key = (name, reason) + count = self.declines.get(key, 0) + 1 + self.declines[key] = count + if count == 1 and self.rank0: + at = f" ({shape})" if shape else "" + logger.warning(f"{TAG} lever {name} declined a step{at}: {reason}") + + def step(self) -> None: + """Count one executor step; every ``DECLINE_REPORT_STEPS`` log the decline counts.""" + self.clock[0] += 1 + if self.clock[0] % DECLINE_REPORT_STEPS == 0 and self.declines and self.rank0: + counts = {f"{name}: {reason}": count for (name, reason), count in self.declines.items()} + logger.info(f"{TAG} lever declines after {self.clock[0]} steps: {counts}") + + def report(self) -> str: + """The installed lever table as one JSON line (``levers.LeverTable.as_json``).""" + if self.levers is None: + return "{}" + return self.levers.as_json({name: ok for name, (ok, _) in self.reports.items()}) + def __bool__(self) -> bool: return bool(self.paths) @@ -1031,101 +1249,101 @@ def dependents(names: Iterable[str]) -> FrozenSet[str]: return frozenset(path for path in FAST_PATH_NAMES if closure((path,)) & names) -def windowed_context_arena(spec_worker: Any) -> bool: - """Whether the DFlash worker keeps only a window of each slot's context K/V. - - Read off the worker the built-in model built (``DFlashWorker.windowed_ctx_arena``). - """ - return bool(getattr(spec_worker, "windowed_ctx_arena", False)) - - -_WINDOWED_ARENA = "a windowed drafter context arena" +#: The levers that write the dense full-length drafter arena by its own column rule, and so +#: need it: the context-K/V side chain with its fused store, its early fc fork, and the +#: folded commit that also writes the drafter's preamble. +DENSE_ARENA_LEVERS: Tuple[str, ...] = ( + "drafter_kv_side_chain", + "drafter_fc_early_fork", + "drafter_preamble_from_commit", +) @dataclasses.dataclass(frozen=True) -class CapacityGate: - """A served fast path left out above a request capacity, because of what it costs. +class StructuralRules: + """What a deployment's drafter arena, architecture and facts exclude or require. - ``path`` (and every fast path built on it) is in the plan only when the - engine's ``max_batch_size`` is at most ``most_requests``; an unknown - ``max_batch_size`` leaves it out. ``cost`` says in words what the path takes - from the deployment. + ``exclusions`` maps each served lever a rule leaves out to the reason; ``requires`` + adds the levers a served lever needs on this deployment (``levers.resolve``). """ - path: str - most_requests: int - cost: str - - -#: The fast paths a plan carries only at a few concurrent requests. -#: ``half_window_experts`` stacks a second half expert window on every TP4 rank, -#: which takes 45 GiB per rank from the KV cache; it is measured ahead of the -#: owner-only window through 8 concurrent requests, and beyond that the -#: prefix cache's capacity is worth more. -CAPACITY_GATES: Tuple[CapacityGate, ...] = ( - CapacityGate( - "half_window_experts", - 8, - "a second half expert window per TP4 rank, 45 GiB of KV cache", - ), -) + exclusions: Mapping[str, str] + requires: Mapping[str, Tuple[str, ...]] -def capacity_gated( - paths: FrozenSet[str], max_batch_size: Optional[int] -) -> Tuple[FrozenSet[str], List[CapacityGate]]: - """``paths`` without the gates closed at ``max_batch_size``, and the gates that closed. +def structural_rules(served: FrozenSet[str], deployment: "_levers.Deployment") -> StructuralRules: + """The structural rules of ``deployment`` for the levers of ``served``. Pure. - A gate closes when its path is in ``paths`` and ``max_batch_size`` is None or - above its ``most_requests``; the path and its dependents leave the set, which - stays closed under ``depends_on``. Changes nothing. + * Over a windowed arena, on an architecture outside ``WINDOWED_COMMIT_SMS``, the + arena-writing paths (``FULL_CONTEXT_ARENA_PATHS``) need this target's compaction + (``WINDOWED_CONTEXT_ARENA_PATHS``), and are left out where it is not served. + * Any other arena leaves the compaction out. + * ``DENSE_ARENA_LEVERS`` need a dense full-length arena. + * ``windowed_commit_parallel`` runs inside the windowed commit of an architecture of + ``WINDOWED_COMMIT_SMS``, so it needs a windowed arena there. + * A paged arena leaves out ``incremental_block_offsets``: the step-input stager + declines the block-table copy while a draft KV cache manager is attached. + * ``TRTLLM_KV_DEFER_CLOSE=1`` leaves out ``step_loop_on_execution_stream``, which is + not validated with the deferred close. """ - closed = [ - gate - for gate in CAPACITY_GATES - if gate.path in paths and (max_batch_size is None or max_batch_size > gate.most_requests) - ] - if not closed: - return paths, closed - return paths - dependents(gate.path for gate in closed), closed + arena = deployment.arena + exclusions: Dict[str, str] = {} + requires: Dict[str, Tuple[str, ...]] = {} + if arena is _levers.Arena.WINDOWED: + if deployment.sm not in WINDOWED_COMMIT_SMS: + compaction = tuple(name for name in WINDOWED_CONTEXT_ARENA_PATHS if name in served) + for name in FULL_CONTEXT_ARENA_PATHS: + if name not in served: + continue + if compaction: + requires[name] = compaction + else: + exclusions[name] = ( + f"a windowed drafter context arena on sm{deployment.sm}, which serves " + "neither the window compaction nor a commit that compacts" + ) + else: + for name in WINDOWED_CONTEXT_ARENA_PATHS: + if name in served: + exclusions[name] = f"the drafter context arena is {arena.value}, not windowed" + if arena is not _levers.Arena.DENSE: + for name in DENSE_ARENA_LEVERS: + if name in served: + exclusions[name] = ( + f"it needs a dense full-length drafter context arena, not a {arena.value} one" + ) + if "windowed_commit_parallel" in served and ( + arena is not _levers.Arena.WINDOWED or deployment.sm not in WINDOWED_COMMIT_SMS + ): + exclusions["windowed_commit_parallel"] = ( + "it needs a windowed drafter context arena on an architecture whose commit " + f"compacts (sm{', sm'.join(str(sm) for sm in sorted(WINDOWED_COMMIT_SMS))}); this " + f"is a {arena.value} arena on sm{deployment.sm}" + ) + if arena is _levers.Arena.PAGED and "incremental_block_offsets" in served: + exclusions["incremental_block_offsets"] = ( + "the step-input stager declines the block-table copy over a paged drafter context" + ) + if deployment.facts.kv_defer_close and "step_loop_on_execution_stream" in served: + exclusions["step_loop_on_execution_stream"] = ( + f"{_levers.KV_DEFER_CLOSE_ENV}=1: the deferred close is not validated with it" + ) + return StructuralRules(exclusions, requires) @dataclasses.dataclass(frozen=True) -class Yield: - """A served fast path left out of a plan that carries another one doing its job. - - ``path`` (and every fast path built on it) leaves the plan when ``to`` is in it, - after the capacity gates. ``why`` says in words what the two share. - """ - - path: str - to: str - why: str +class _EngineFacts: + """The engine facts ``compose`` falls back to without ``_router_index.DeploymentFacts``.""" - -#: The fast paths a plan leaves to another one it carries. ``hot_replicas`` and -#: ``half_window_experts`` both stack extra NVFP4 experts on every TP4 rank and place -#: each step's experts over them; the half window is measured ahead through 8 -#: concurrent requests (its capacity gate), so a plan carrying it leaves the hot copies -#: out and the two never both hold memory or both arm. -YIELDS: Tuple[Yield, ...] = ( - Yield( - "hot_replicas", - "half_window_experts", - "both stack extra NVFP4 experts per TP4 rank and place each step's experts over them", - ), -) + max_batch_size: Optional[int] + max_num_tokens: Optional[int] = None + disaggregated: Optional[bool] = None -def yielded(paths: FrozenSet[str]) -> Tuple[FrozenSet[str], List[Yield]]: - """``paths`` without those that yield to another path in it, and the yields applied. - - Closed under ``depends_on`` like ``capacity_gated``. Changes nothing. - """ - applied = [rule for rule in YIELDS if rule.path in paths and rule.to in paths] - if not applied: - return paths, applied - return paths - dependents(rule.path for rule in applied), applied +#: The prefix of the drafter arena item of the TP agreement. +_ARENA_TAG = "drafter context arena: " +#: The prefix of a rank's lever error carried through the TP agreement. +_LEVER_ERROR = "lever table: " def compose( @@ -1136,81 +1354,149 @@ def compose( compute_capability: Optional[Tuple[int, int]], spec_worker: Any = None, max_batch_size: Optional[int] = None, + *, + draft_model: Any = None, + facts: Any = None, + snapshot: Optional[Mapping[str, str]] = None, ) -> Plan: - """The plan this model is built with: its composition, or ``NO_PATHS``. + """The plan this model is built with: its composition's lever table, or ``NO_PATHS``. Call after the built-in constructor and before anything is installed. A - deployment without a composition (``composition``, structural) gets - ``NO_PATHS`` under ``auto`` and raises under ``require``. Otherwise the plan - is the composition's fast paths, less those written for a windowed arena when - the worker's arena is not windowed, and, over a windowed arena, less the - arena-writing paths (``FULL_CONTEXT_ARENA_PATHS``) unless the composition also - serves the window compaction or the architecture is one of - ``WINDOWED_COMMIT_SMS`` (whose windowed commit addresses the arena by its - physical columns and compacts itself; structural as well), less the paths a - capacity gate closes at the engine's ``max_batch_size`` (``CAPACITY_GATES``; - every rank reads the same engine argument), less the paths that yield to another - path still in the plan (``YIELDS``), and the preconditions of those - paths and the arena kind are checked once, agreed over the TP group - (collective; every rank returns the same plan or raises the same error). + deployment without a composition (``composition``, structural) gets ``NO_PATHS`` + under ``auto`` and raises under ``require``; its lever environment is still + validated (``levers.no_composition_notes``). Otherwise: + + 1. the drafter context arena (``levers.arena_of``, the worker's own predicate over + ``spec_worker`` and ``draft_model``) and the ``levers.Deployment`` + (``levers.deployment_of``; ``facts`` are the engine's + ``_router_index.DeploymentFacts``, else ``max_batch_size`` alone); + 2. the lever table (``levers.resolve`` under ``structural_rules``), whose levers on + are the plan; + 3. the preconditions of those levers, the arena, the lever environment and any + rank's lever error, agreed over the TP group in one collective: every rank + returns the same plan or raises the same error. + + The lever table is logged once, on TP rank 0. ``snapshot`` stands in for the + process's environment snapshot (``levers.environment``). Raises: ValueError: under ``modeling_v2="require"`` when the deployment has no - composition; under any mode when a precondition is missing on some - rank, or the ranks disagree on the drafter arena (on every rank alike). + composition; under any mode for a lever error (``levers.resolve``), a + precondition missing on some rank, or ranks that disagree on the drafter + arena or the lever environment (on every rank alike). """ + snapshot = _levers.environment() if snapshot is None else snapshot served = composition(model_config, compositions, compute_capability) mapping = getattr(model_config, "mapping", None) if served is None: require_served(model_config, mode, compositions, compute_capability) + notes = _levers.no_composition_notes(snapshot, DEPENDS_ON) reason = unserved_reason(model_config, compositions, compute_capability) logger.info( f"{TAG} no fast path for this deployment ({reason}); running the built-in forward" ) + for note in notes: + logger.info(f"{TAG} {note}") return NO_PATHS from tensorrt_llm._torch.speculative.dflash import _PAGED_ATTENTION_BACKENDS - # The worker pages the drafter's context on these backends, and windows only a dense arena. - paged = getattr(spec_worker, "_dflash_attention_backend", None) in _PAGED_ATTENTION_BACKENDS - windowed = windowed_context_arena(spec_worker) and not paged - paths, gated = capacity_gated(frozenset(served.served), max_batch_size) - paths, yields = yielded(paths) - if windowed: - if not paths & frozenset(WINDOWED_CONTEXT_ARENA_PATHS) and ( - sm_of(compute_capability) not in WINDOWED_COMMIT_SMS - ): - paths = paths - dependents(FULL_CONTEXT_ARENA_PATHS) - else: - paths = paths - dependents(WINDOWED_CONTEXT_ARENA_PATHS) - local = missing_preconditions(text_model, model_config, preconditions(paths, served)) - agreed = agreed_union(mapping, local + ([_WINDOWED_ARENA] if windowed else [])) - missing = [item for item in agreed if not item.startswith(_WINDOWED_ARENA)] - split = [ - item for item in agreed if item.startswith(_WINDOWED_ARENA) and item != _WINDOWED_ARENA - ] - if split: - missing.append(f"the same drafter context arena on every rank ({split[0]})") + arena = _levers.arena_of(spec_worker, draft_model, _PAGED_ATTENTION_BACKENDS) + engine = facts if facts is not None else _EngineFacts(max_batch_size) + table: Optional[_levers.LeverTable] = None + local: List[str] = [] + try: + # Every grammar error first, the role hint's included, before the role is derived. + _levers.parse_knobs(snapshot, DEPENDS_ON) + deployment = _levers.deployment_of( + model_config, compute_capability, arena, engine, snapshot + ) + rules = structural_rules(served.served, deployment) + table = _levers.resolve( + deployment, served.served, rules.exclusions, snapshot, DEPENDS_ON, rules.requires + ) + except ValueError as error: + local.append(f"{_LEVER_ERROR}{error}") + if table is not None: + local += missing_preconditions( + text_model, model_config, preconditions(table.paths(), served) + ) + tags = [f"{_ARENA_TAG}{arena.value}", *_levers.tagged(snapshot)] + agreed = agreed_union(mapping, local + tags) + own = set(tags) + tagged = (_ARENA_TAG, _levers.AGREEMENT_TAG) + split = [item for item in agreed if item.startswith(tagged) and item not in own] + errors = [item[len(_LEVER_ERROR) :] for item in agreed if item.startswith(_LEVER_ERROR)] + missing = [item for item in agreed if not item.startswith(tagged + (_LEVER_ERROR,))] composed = ( f"{TAG} sm{sm_of(compute_capability)} tep{mapping.tp_size} on {served.topology.description}" ) + if split: + raise ValueError( + f"{composed}: the TP ranks disagree on the drafter context arena or the lever " + "environment: " + "; ".join(split) + ) + if errors: + raise ValueError(f"{composed}: " + "; ".join(errors)) if missing: raise ValueError(f"{composed}: the fast paths need " + "; ".join(missing)) - plan = Plan(paths) - if windowed: - composed += f" with {_WINDOWED_ARENA}" - if paged: - composed += " with a paged drafter context" - for gate in gated: - composed += ( - f" without {gate.path} (max_batch_size {max_batch_size} above {gate.most_requests};" - f" it costs {gate.cost})" - ) - for rule in yields: - composed += f" without {rule.path} (it yields to {rule.to}: {rule.why})" - logger.info(f"{composed}: {plan.describe()}") + rank0 = int(getattr(mapping, "tp_rank", 0) or 0) == 0 + plan = Plan(table.paths(), levers=table, rank0=rank0) + logger.info(f"{composed} with a {arena.value} drafter context arena: {plan.describe()}") + if rank0: + for line in table.describe(): + logger.info(f"{TAG} {line}") + for warning in table.warnings: + logger.warning(f"{TAG} lever warning: {warning}") return plan +def table_for( + compositions: Mapping[CompositionKey, TargetPaths], + sm: int, + tp: int, + max_batch_size: Optional[int], + role: str, + arena: str, + knobs: Optional[Mapping[str, str]] = None, + facts: Optional["_levers.Facts"] = None, + max_num_tokens: Optional[int] = None, + role_source: Optional[str] = None, +) -> "_levers.LeverTable": + """The lever table ``compose`` resolves for a deployment, without a model or a device. + + For tests and for reviewing a deployment's levers. ``role`` and ``arena`` are ``levers.Role`` / + ``levers.Arena`` values; ``knobs`` the lever environment (names under + ``levers.PREFIX``); ``facts`` the shared-path facts; ``role_source`` defaults to what + a served deployment of that role reports (agg servers derive their role, the others + take the hint). + + Raises: + KeyError: for a key no composition serves. + ValueError: as ``levers.resolve``. + """ + served = compositions[(sm, tp)].served + role = _levers.Role(role) + if role_source is None: + role_source = { + _levers.Role.AGG: "no cache transceiver", + _levers.Role.UNKNOWN: _levers.DISAGGREGATED_UNHINTED, + }.get(role, "hint") + deployment = _levers.Deployment( + sm=sm, + tp=tp, + max_batch_size=max_batch_size, + max_num_tokens=max_num_tokens, + role=role, + role_source=role_source, + arena=_levers.Arena(arena), + facts=facts or _levers.Facts(), + ) + rules = structural_rules(served, deployment) + return _levers.resolve( + deployment, served, rules.exclusions, dict(knobs or {}), DEPENDS_ON, rules.requires + ) + + # -------------------------------------------------------------------------- # Install refusals agreed over the TP group # -------------------------------------------------------------------------- diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py index f159f2459851..a5bcf39d14aa 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/drafter.py @@ -58,14 +58,28 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None: # verify_fold then composes on top of it. from .kernels.context_kv_fork import fork as _context_kv_fork - _context_kv_fork.install(shell.llm, shell.llm.model) + parallel = plan.enabled("windowed_commit_parallel") + _context_kv_fork.install( + shell.llm, + shell.llm.model, + side_chain=plan.enabled("drafter_kv_side_chain"), + early_fc=plan.enabled("drafter_fc_early_fork"), + parallel_min_requests=( + plan.param("windowed_commit_parallel", "MIN_REQUESTS") if parallel else None + ), + declined=plan.declined, + ) verify_fold = plan.enabled("verify_fold") path_walk_tail = plan.enabled("path_walk_tail") if verify_fold or path_walk_tail: from .kernels.verify_fold import fold as _spec_loop_fold _spec_loop_fold.install( - getattr(shell.llm, "spec_worker", None), verify_fold, path_walk_tail + getattr(shell.llm, "spec_worker", None), + verify_fold, + path_walk_tail, + preamble=plan.enabled("drafter_preamble_from_commit"), + declined=plan.declined, ) if plan.enabled("drafter_hops"): from .kernels.drafter_hops import hops as _drafter_hops @@ -84,12 +98,21 @@ def adopt_drafter(shell, plan: "_dispatch.Plan") -> None: if plan.enabled("split_sampler"): from .kernels.split_sampler import target_head as _target_head - _target_head.install( + multimem = plan.enabled("head_multimem_gather") + coop = plan.enabled("head_sampler_coop_fold") + head = _target_head.install( shell.llm, shell._mapping, plan.enabled("lm_head_inbox"), plan.enabled("vocab_parallel_head"), + multimem_gather=multimem, + coop_max_rows=plan.param("head_sampler_coop_fold", "MAX_ROWS") if coop else 0, ) + if multimem: + # The gather's refusal is agreed over the TP group, so every rank reports alike. + plan.installed( + "head_multimem_gather", head.multimem is not None, str(head.multimem_refusal) + ) if plan.enabled("wide_split_sampler"): _target_head.widen(shell.llm) # Last: the arena compaction derives its worker class from whatever class every diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py index bdedd1c71ca7..8b1dd6067abe 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/layers.py @@ -113,10 +113,17 @@ class DirectPublish: ``PostAttnHopCall.armed`` calls ``recover`` on any exception after ``arm``: a collective reset of the transport. A reset that fails poisons this object: every later ``arm`` raises. + + ``rows`` are the row counts this publication serves, each on the same workspace: by + default ``ROWS``; the wide direct publication (``wide_oproj_publish``) serves + ``WIDE_PUBLISH_ROWS`` on the 32-row direct transport. A step of any of them arms every + layer with that row count's slots, so the parity sequence above holds across steps of + different row counts too. """ - #: The one row count published directly: the K = 7 verify step. Other row counts keep - #: the post-attention hop, on its own workspace. + #: The one row count the 8-row publication serves: the K = 7 verify step. Other row + #: counts keep the post-attention hop, on its own workspace, unless the wide + #: publication serves them. ROWS = 8 #: The o_proj kernel families that store straight into the slot, per device compute #: capability: the families of that arch's ``post_attn_hop.VALIDATED_PRODUCERS``. The @@ -131,6 +138,7 @@ def __init__( workspace: "_hop.PostAttnDirectWorkspace", dispatchers: tuple, joins: "tuple[MuseSparkOpenTepPostAttnJoin, ...]", + rows: Tuple[int, ...] = (ROWS,), ) -> None: if len(dispatchers) != len(joins) or len(joins) % 2: raise ValueError( @@ -142,11 +150,23 @@ def __init__( self.workspace = workspace self.dispatchers = tuple(dispatchers) self.joins = tuple(joins) - self.slots = tuple(workspace.slot(index % 2, self.ROWS) for index in range(len(joins))) + #: The row counts this publication serves; ``slots`` are those of the first. + self.rows = tuple(int(r) for r in rows) + self.slots_by_rows = { + r: tuple(workspace.slot(index % 2, r) for index in range(len(joins))) for r in self.rows + } + self.slots = self.slots_by_rows[self.rows[0]] #: Why the transport could not be reset after a failed step, or ``None``. self.poisoned: Optional[str] = None - def check_step(self, residual: torch.Tensor) -> None: + def serves(self, rows: Optional[int]) -> bool: + """Whether a decode step of ``rows`` rows publishes through this object.""" + return rows in self.slots_by_rows + + def _label(self) -> str: + return "/".join(str(r) for r in self.rows) + + def check_step(self, residual: torch.Tensor, rows: Optional[int] = None) -> None: """Before anything is published: the step's input residual is what the consumer takes. Every layer's consumer call reads the residual the previous layer returned, which @@ -156,26 +176,30 @@ def check_step(self, residual: torch.Tensor) -> None: left to check, and it is checked before layer 0's o_proj stores anything. Raises: - ValueError: if ``residual`` is not a contiguous FP32 ``[ROWS, 8192]`` tensor on - the workspace's device, or the step is being compiled. + ValueError: if ``residual`` is not a contiguous FP32 ``[rows, 8192]`` tensor on + the workspace's device, or the step is being compiled. ``rows`` is the step's + row count (by default the first this object serves). """ + rows = self.rows[0] if rows is None else int(rows) device = getattr(self.workspace, "device", residual.device) if ( residual.dtype != torch.float32 - or tuple(residual.shape) != (self.ROWS, _hop.HIDDEN) + or tuple(residual.shape) != (rows, _hop.HIDDEN) or not residual.is_contiguous() or residual.device != device or torch.compiler.is_compiling() ): raise ValueError( - f"[musespark_open] attention_output_into_hop: a {self.ROWS}-row step's residual is {residual.dtype} " + f"[musespark_open] attention_output_into_hop: a {rows}-row step's residual is {residual.dtype} " f"{tuple(residual.shape)} on {residual.device} (contiguous=" f"{residual.is_contiguous()}); the consumer hop takes a contiguous FP32 " - f"[{self.ROWS}, {_hop.HIDDEN}] tensor on {device}, not compiled" + f"[{rows}, {_hop.HIDDEN}] tensor on {device}, not compiled" ) - def arm(self) -> None: - """Hand every layer's o_proj its slot and every join its slot and parity. + def arm(self, rows: Optional[int] = None) -> None: + """Hand every layer's o_proj its slot and every join its slot, parity and this object. + + ``rows`` is the step's row count, one of ``self.rows`` (by default the first). Raises: RuntimeError: if an earlier failed step left the transport unreset. @@ -185,11 +209,12 @@ def arm(self) -> None: "[musespark_open] attention_output_into_hop: the direct transport was not reset after " "a failed step " f"({self.poisoned}); its buffers may hold words no call consumed, so no " - f"{self.ROWS}-row step may publish into it" + f"{self._label()}-row step may publish into it" ) + slots = self.slots_by_rows[self.rows[0] if rows is None else int(rows)] for index, (dispatcher, join) in enumerate(zip(self.dispatchers, self.joins)): - dispatcher.publish = self.slots[index] - join._direct = (self.slots[index], index % 2) + dispatcher.publish = slots[index] + join._direct = (slots[index], index % 2, self) def disarm(self) -> None: for dispatcher, join in zip(self.dispatchers, self.joins): @@ -210,7 +235,7 @@ def recover(self, error: BaseException) -> None: except RuntimeError: pass logger.error( - f"[musespark_open] attention_output_into_hop: a {self.ROWS}-row decode step raised while " + f"[musespark_open] attention_output_into_hop: a {self._label()}-row decode step raised while " "o_proj published into the " f"post-attention hop's direct transport ({error!r}); resetting the transport on " "every rank of the group (collective: a rank that did not raise leaves this one " @@ -223,7 +248,7 @@ def recover(self, error: BaseException) -> None: logger.error( "[musespark_open] attention_output_into_hop: the direct transport could not be reset " f"({reset_error!r}); every " - f"later {self.ROWS}-row step raises" + f"later {self._label()}-row step raises" ) else: logger.warning( @@ -237,7 +262,9 @@ class PostAttnHopCall: With ``attention_output_into_hop``, also the ``DirectPublish`` that steps of its row count arm. With ``wide_post_attention_hop``, also the wide transport (``wide``) that steps of - ``MAX_ROWS + 1 .. WIDE_MAX_ROWS`` rows run on. + ``MAX_ROWS + 1 .. WIDE_MAX_ROWS`` rows run on. With both and ``wide_oproj_publish``, + also the wide ``DirectPublish`` (``wide_direct``) that steps of + ``WIDE_PUBLISH_ROWS`` rows arm instead of the wide transport. """ def __init__( @@ -246,12 +273,23 @@ def __init__( joins: "tuple[MuseSparkOpenTepPostAttnJoin, ...]" = (), direct: Optional[DirectPublish] = None, wide: "Optional[_hop.PostAttnHopWorkspace]" = None, + wide_direct: Optional[DirectPublish] = None, ) -> None: self.op = _hop.post_attn_hop self.workspace = workspace self.joins = tuple(joins) self.direct = direct self.wide = wide + self.wide_direct = wide_direct + + def publication_for(self, rows: Optional[int]) -> Optional[DirectPublish]: + """The direct publication a decode step of ``rows`` rows arms, or ``None`` (the hop).""" + if self.direct is not None and rows == DirectPublish.ROWS: + return self.direct + wide_direct = getattr(self, "wide_direct", None) + if wide_direct is not None and wide_direct.serves(rows): + return wide_direct + return None def takes_call( self, @@ -294,20 +332,21 @@ def armed(self, rows: Optional[int] = None, residual: Optional[torch.Tensor] = N """Arm every join for the duration of one decode forward; disarm on any exit. A step of ``DirectPublish.ROWS`` rows (``attention_output_into_hop``) also arms the - direct publication, + direct publication, and a step of ``WIDE_PUBLISH_ROWS`` rows the wide one when it is + built (``publication_for``), after ``DirectPublish.check_step`` on the step's input ``residual`` (when given); an exception after that hands the transport to ``DirectPublish.recover`` before it propagates. """ - direct = self.direct if self.direct is not None and rows == DirectPublish.ROWS else None + direct = self.publication_for(rows) for join in self.joins: join._armed = self direct_armed = False try: if direct is not None: if residual is not None: - direct.check_step(residual) - direct.arm() + direct.check_step(residual, rows) + direct.arm(rows) direct_armed = True yield except Exception as error: @@ -453,7 +492,8 @@ class MuseSparkOpenTepPostAttnJoin(MuseSparkOpenGatedResidual): kernel cannot read in place is copied (``readable``) rather than declined. With ``attention_output_into_hop`` armed for the step (``DirectPublish.arm``) the join also holds - ``_direct``, its layer's slot and parity: o_proj has already written the + ``_direct``, its layer's slot, parity and the ``DirectPublish`` that armed it (the 8-row + one or the wide one, whose op and workspace the call runs on): o_proj has already written the partial into that slot of every rank, so the join runs the consumer hop. That call cannot fall back to the built-in all-reduce (every peer polls this rank's slot, and nothing may read the slot's multicast address), so a call @@ -523,15 +563,16 @@ def _add_after_rmsnorm_and_norm(self, residual, branch, eps, next_norm, gain=Non def _consume_published(self, call, residual, branch, eps, next_norm, gain): """The consumer hop on the slot o_proj published into (see the class docstring).""" - slot, parity = self._direct + slot, parity, direct = self._direct alpha, beta = self.coefficients() weight = next_norm.weight + max_rows = getattr(direct.workspace, "max_rows", _hop.MAX_ROWS) if not ( gain is None and self._fused_enabled and next_norm.use_gemma and not torch.compiler.is_compiling() - and _hop.hop_takes(slot, residual, alpha, beta, weight) + and _hop.hop_takes(slot, residual, alpha, beta, weight, max_rows) and tuple(branch.shape) == tuple(slot.shape) ): raise RuntimeError( @@ -549,7 +590,6 @@ def _consume_published(self, call, residual, branch, eps, next_norm, gain): ) _hop.publish_copy(slot, branch) residual, alpha, beta, weight = _hop.readable(residual, alpha, beta, weight) - direct = call.direct return direct.op( slot, residual, @@ -853,6 +893,123 @@ def _build_direct_publish( return DirectPublish(workspace, tuple(dispatchers), call.joins) +#: The row counts the wide direct publication serves: two to four requests' K = 7 verify steps. +WIDE_PUBLISH_ROWS = (16, 24, 32) +#: Whether the published wide row counts fork the o_proj -> pool-weight prefetch: as at 8 rows, +#: the consumer hop ends before the prefetch would, so none is forked. +WIDE_PUBLISH_DROPS_PREFETCH = True + + +def _build_wide_direct_publish( + model: MuseSparkOpenModel, call: PostAttnHopCall, early_trigger: bool = False +) -> Tuple[Optional[DirectPublish], Optional[str]]: + """``attention_output_into_hop`` on the wide TP4 steps (``wide_oproj_publish``). + + On a decode step of ``WIDE_PUBLISH_ROWS`` rows layer ``l``'s o_proj (the ``wide_dense_gemms`` + split-K tactic of the step's rows) stores straight into this rank's ``[rows, 8192]`` slot of + a 32-row direct transport (``PostAttnDirectWorkspace(max_rows=WIDE_MAX_ROWS)``, buffer + ``l mod 2``), and the join runs the consumer hop there instead of the wide hop + (``wide_post_attention_hop``), which would load the partial and publish it itself. Same + bits: the same tactic writes the same partial; the consumer's rank-order sum and 4-CTA + boundary are the wide hop's statements (a published ``-0`` adds what the hop's cleared + ``+0`` adds). + + Collective: every rank of the TP group calls it when the plan carries the lever. Builds + on the 8-row publication (``call.direct``: its o_proj dispatchers and their producer + policy) and the wide hop (``call.wide``: the joins take 9..32-row calls); checked first + and agreed over the group before the workspace's collective construction: TP4, and an + o_proj kernel of the policy's families at every served row count on every layer. A + refusal leaves every rank on the wide hop. Returns ``(the wide publication, None)`` or + ``(None, the agreed refusal)``. + """ + from .ops.gemm.dense_gemm import parse_kernel + + layers = list(model.layers) + mapping = layers[0].mapping + direct = call.direct + reasons = [] + if direct is None: + reasons.append("o_proj does not publish its 8-row steps (attention_output_into_hop)") + if call.wide is None: + reasons.append("no wide post-attention hop transport (wide_post_attention_hop)") + tp_size = getattr(mapping, "tp_size", None) + if tp_size != 4: + reasons.append(f"the wide direct transport is TP4 only, got TP{tp_size}") + kernels = {} + if direct is not None: + for dispatcher in direct.dispatchers: + policy = dispatcher.publish_policy + plan = dispatcher.site.plan + for rows in WIDE_PUBLISH_ROWS: + kernel = plan.kernel_for(dispatcher.site.role, rows) + family = None if kernel is None else parse_kernel(kernel)[0] + if policy is None or family not in policy.direct: + reasons.append( + f"layer {dispatcher.site.layer_idx}: o_proj runs {kernel} at {rows} rows, " + "not a kernel family that publishes directly " + f"({sorted(policy.direct) if policy is not None else []})" + ) + kernels[rows] = kernel + if reasons: + break + refusal = _tp_agreement.agreed_refusal(mapping, " | ".join(reasons) if reasons else None) + if refusal is not None: + return None, f"{refusal}; the wide steps keep wide_post_attention_hop" + variants = sorted( + { + ( + float(layer.post_attn_norm.variance_epsilon), + float(layer.pre_ffn_norm.variance_epsilon), + bool(layer.add_attn._preqkv_norm_fp32), + ) + for layer in layers + } + ) + workspace = _hop.shared_direct_workspace( + mapping, variants=variants, early_trigger=early_trigger, max_rows=_hop.WIDE_MAX_ROWS + ) + dropped = 0 + if WIDE_PUBLISH_DROPS_PREFETCH: + for dispatcher in direct.dispatchers: + site = dispatcher.site + site.skip_prefetch_rows = frozenset(site.skip_prefetch_rows | set(WIDE_PUBLISH_ROWS)) + dropped += int(bool(site.plan.prefetch) and site.next_weight is not None) + described = ", ".join(f"{kernels[rows]}@{rows}" for rows in WIDE_PUBLISH_ROWS) + _hop.log_once_wide_direct(len(layers), WIDE_PUBLISH_ROWS, described, dropped) + return DirectPublish(workspace, direct.dispatchers, call.joins, rows=WIDE_PUBLISH_ROWS), None + + +def _check_wide_publication(direct: DirectPublish, o_proj, rows: int) -> str: + """Layer 0's real dispatcher publishes a ``rows``-row product into the wide transport's + slot 0 inside a capture that is never replayed; the workspace checks the captured kernel + (``check_producer``, collective). Returns its name.""" + dispatcher = direct.dispatchers[0] + slot = direct.slots_by_rows[rows][0] + weight = o_proj.weight + x = torch.zeros(rows, int(weight.shape[1]), dtype=weight.dtype, device=weight.device) + site = dispatcher.site + + def publish() -> None: + dispatcher.publish = slot + # The capture holds the producer only: no prefetch fork at the checked row count. + skipped = site.skip_prefetch_rows + site.skip_prefetch_rows = frozenset(skipped | {rows}) + try: + with torch.inference_mode(): + dispatcher.apply(o_proj, x, weight, None) + finally: + dispatcher.publish = None + site.skip_prefetch_rows = skipped + + def warm() -> None: + # The same kernel, disarmed, into a fresh output: library state and a flashinfer + # kernel's compilation happen outside the capture. + with torch.inference_mode(): + dispatcher.apply(o_proj, x, weight, None) + + return direct.workspace.check_producer(publish, warm) + + def verify_direct_publish(text_model: MuseSparkOpenModel) -> str: """``attention_output_into_hop``, once the weights are final: the kernel the published o_proj runs. @@ -897,7 +1054,13 @@ def warm() -> None: with torch.inference_mode(): dispatcher.apply(o_proj, x, weight, None) - return direct.workspace.check_producer(publish, warm) + name = direct.workspace.check_producer(publish, warm) + wide = getattr(call, "wide_direct", None) + if wide is not None: + # The wide publication's tactics (one per served row count), the same check. + for rows in wide.rows: + _check_wide_publication(wide, o_proj, rows) + return name def _raise_refusal(reason: Optional[str], path: str) -> None: @@ -1092,8 +1255,15 @@ def adopt_model( from .ops.attention.swa_block import MULTI_MAX_REQUESTS as _swa_multi_max min_seqs = 1 + max(_full_multi_max, _swa_multi_max) + # packed_prep_append: calls of its MIN_ROWS rows or more on the packed kernel. + packed_min_rows = ( + plan.param("packed_prep_append", "MIN_ROWS") + if plan.enabled("packed_prep_append") + else None + ) _tp_agreement.agreed_install( - model.layers[0].mapping, lambda: multi_request_append.install(model, min_seqs) + model.layers[0].mapping, + lambda: multi_request_append.install(model, min_seqs, packed_min_rows), ) if decode_kv_splits is not None: # After every other attention re-class: the mixed-step mixin goes in front of all of @@ -1152,7 +1322,11 @@ def adopt_model( if plan.enabled("wide_prologue"): from .kernels.wide_prologue import prologue as wide_prologue - wide_prologue.install(model) + wide_prologue.install( + model, + router_rows=plan.enabled("wide_router_per_row"), + direct=plan.enabled("wide_direct_transports"), + ) router_gemm = None if plan.enabled("router_gemm"): # After moe_prologue and the router fork, which check the MoE class they re-class: @@ -1243,6 +1417,22 @@ def adopt_model( from .kernels.pool_into_route import route as _pool_direct pool_direct = _pool_direct.install(model, dense_gemm_plan, mapping=model.layers[0].mapping) + if plan.enabled("wide_direct_transports"): + # Both transports' wide layouts follow construction-time state wide_prologue.install + # agreed over the TP group, so every rank reports the same. + unbuilt = [ + name + for name, built in ( + ("pool_into_route", pool_direct), + ("unpool_into_boundary", unpool_direct), + ) + if built is None or built.wide_workspace is None + ] + plan.installed( + "wide_direct_transports", + not unbuilt, + "no wide direct transport on " + ", ".join(unbuilt), + ) # finalize_transport replaces only the object every MoE layer calls its fused # finalize on, whatever class the MoE has. finalize_transport = None @@ -1251,7 +1441,9 @@ def adopt_model( # and runs nothing beyond the built-in modules. from .kernels.finalize_transport import transport as _finalize_transport - finalize_transport = _finalize_transport.install(model) + finalize_transport = _finalize_transport.install( + model, wide=plan.enabled("wide_finalize_transport") + ) wide_moe_allreduce = None if plan.enabled("wide_moe_allreduce"): # Replaces only the module every MoE layer reduces its unfused output with; @@ -1280,6 +1472,12 @@ def adopt_model( if plan.enabled("attention_output_into_hop"): # After dense_gemms (its o_proj dispatchers) and post_attention_hop (its joins). post_attn_hop.direct = _build_direct_publish(model, post_attn_hop, early_trigger=True) + if plan.enabled("wide_oproj_publish"): + # The wide steps' publication (a collective agreement on every rank). + post_attn_hop.wide_direct, refusal = _build_wide_direct_publish( + model, post_attn_hop, early_trigger=True + ) + plan.installed("wide_oproj_publish", refusal is None, refusal or "") wide_hop = None if plan.enabled("multi_request_post_attention_hop"): # After post_attention_hop, whose joins it arms on 9..1024-row decode steps. diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py new file mode 100644 index 000000000000..0a9acf7af56c --- /dev/null +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py @@ -0,0 +1,1138 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The lever table: which installable pieces of this target a model is built with, and why. + +A lever is any installable piece of the target: every entry of ``dispatch.FAST_PATHS``, +the fast paths themselves and the refinements and installs registered beside them. Once +per process ``dispatch.compose`` resolves one ``LeverTable`` from the composition's served +set, the ``Deployment`` (the GPU, the tensor-parallel size, the engine's +``max_batch_size`` and ``max_num_tokens``, the server's role and the drafter's context +arena), the structural rules ``dispatch`` derives from it, and the environment snapshot +(``environment``). The installers then ask the ``dispatch.Plan`` built from the table. + +Knobs. These are the only variables the target reads for levers:: + + MUSESPARK_OPEN_LEVER_=auto|on|off one per lever; unset or empty: the default + MUSESPARK_OPEN_LEVER___= a parameter the lever declares (Rule.params) + MUSESPARK_OPEN_LEVER__DEFAULT=auto|baseline what an unset lever knob means + MUSESPARK_OPEN_LEVER__ROLE=auto|agg|ctx|gen the server's role; unset or auto: derived + +```` is the lever's name upper-cased; ``__`` introduces a parameter or a resolver +setting and never occurs inside a lever name. Values are case-insensitive and +whitespace-stripped. Two shared-path variables are read as facts and never set or retired +here: ``TRTLLM_KV_DEFER_CLOSE`` (on iff exactly ``1``, the KV cache manager's own test) and +``UCX_RNDV_THRESH`` (set or not). + +Semantics, in the order a lever is resolved: + +* structural: a fact the lever's code needs (an arena kind, an architecture, a fact + variable; ``dispatch.structural_rules``) excludes it; nothing overrides that; +* dependency: a lever it builds on (``dispatch.FastPath.depends_on``) or requires + (``Rule.requires``) is off; +* knob: ``off`` removes it, and with it every lever built on it or requiring it; ``on`` + keeps it whatever its deployment rule and gates say; +* default: ``auto`` runs the lever's deployment rule (``Rule.auto``), and a lossless-numeric + lever also needs one of its accuracy gates (``Rule.gates``); ``baseline`` turns every + lever off that is not the upstream target's, and gives the upstream ones their upstream + rule (``Rule.baseline``). + +``on`` for a lever the composition does not serve, that a structural rule excludes, whose +dependency or requirement is off, or that conflicts with another lever that is on, raises. +So does every grammar error and every retired switch (``RETIRED``); ``UCX_RNDV_THRESH`` on +a server with a cache transceiver (a ``ctx`` or ``gen`` role, or an unhinted one); and +``context_fp8_dense`` on where a decode step can reach its ``MIN_ROWS``. The resolver +reads no device and no module of this target: everything it needs is an argument, so it +runs on any host. +""" + +from __future__ import annotations + +import dataclasses +import difflib +import enum +import json +import os +import types +from typing import Callable, Dict, FrozenSet, Iterable, List, Mapping, Optional, Sequence, Tuple + +#: The draft tokens per step of the one speculative-decoding mode the compositions serve. +DRAFT_LEN = 7 +#: Decode rows one request brings to a step: its last accepted token and its drafts. +ROWS_PER_REQUEST = DRAFT_LEN + 1 +#: The fewest decode rows of a step of two requests, where the wide TP4 paths start. +WIDE_MIN_ROWS = 2 * ROWS_PER_REQUEST + +#: The prefix of every lever knob. +PREFIX = "MUSESPARK_OPEN_LEVER_" +#: Separates a lever name from a parameter key. +SEP = "__" +#: What an unset lever knob means: ``auto`` (the default) or ``baseline``. +DEFAULT_KNOB = PREFIX + "_DEFAULT" +#: The server's role: ``auto`` (derived, the default), ``agg``, ``ctx`` or ``gen``. +ROLE_KNOB = PREFIX + "_ROLE" +#: The shared-path facts the resolver reads. +KV_DEFER_CLOSE_ENV = "TRTLLM_KV_DEFER_CLOSE" +UCX_RNDV_THRESH_ENV = "UCX_RNDV_THRESH" + +AUTO = "auto" +ON = "on" +OFF = "off" +BASELINE = "baseline" +LEVER_VALUES = (AUTO, ON, OFF) +DEFAULT_VALUES = (AUTO, BASELINE) + +#: The decision sources, in the order a lever is resolved. +SOURCES = ("structural", "dependency", "knob", "baseline", "auto", "gate") + +#: Retired switches and what replaced them. Setting one, to any value, raises. +RETIRED: Mapping[str, str] = types.MappingProxyType( + { + "MUSESPARK_OPEN_PREP_PACKED_ROWS": PREFIX + "PACKED_PREP_APPEND", + "MUSESPARK_OPEN_CTX_COMMIT_TUNED": PREFIX + "WINDOWED_COMMIT_PARALLEL", + "MUSESPARK_OPEN_CONTEXT_FP8_DENSE": PREFIX + "CONTEXT_FP8_DENSE", + "MUSESPARK_OPEN_CONTEXT_FP8_DENSE_MIN_ROWS": PREFIX + "CONTEXT_FP8_DENSE__MIN_ROWS", + "MUSESPARK_OPEN_CONTEXT_FP8_DENSE_PROJECTIONS": PREFIX + "CONTEXT_FP8_DENSE__PROJECTIONS", + "MUSESPARK_OPEN_CONTEXT_FP8_DENSE_GATHER": PREFIX + "CONTEXT_FP8_ATTN_GATHER", + "TRTLLM_MUSESPARK_TEP_WIDE_DIRECT": PREFIX + "WIDE_DIRECT_TRANSPORTS", + "TRTLLM_MUSESPARK_TEP_WIDE_ROUTER_ROWS": PREFIX + "WIDE_ROUTER_PER_ROW", + "TRTLLM_MUSESPARK_TEP_WIDE_OPROJ_PUBLISH": PREFIX + "WIDE_OPROJ_PUBLISH", + "MUSESPARK_OPEN_STEP_LOOP_ONE_STREAM": PREFIX + "STEP_LOOP_ON_EXECUTION_STREAM", + "MUSESPARK_OPEN_INCREMENTAL_BLOCK_OFFSETS": PREFIX + "INCREMENTAL_BLOCK_OFFSETS", + "MUSESPARK_OPEN_PS_SIDE_CHAIN": PREFIX + "DRAFTER_KV_SIDE_CHAIN", + "MUSESPARK_OPEN_PS_EARLY_FC": PREFIX + "DRAFTER_FC_EARLY_FORK", + "MUSESPARK_OPEN_PS_DRAFTER_PREAMBLE": PREFIX + "DRAFTER_PREAMBLE_FROM_COMMIT", + "MUSESPARK_OPEN_PS_COOP_FOLD": PREFIX + "HEAD_SAMPLER_COOP_FOLD", + "MUSESPARK_OPEN_PS_SYMM_GATHER": PREFIX + "HEAD_MULTIMEM_GATHER", + "MUSESPARK_OPEN_PS_NO_START_PREFETCH": PREFIX + "GRAPH_START_QKV_PREFETCH=off", + "MUSESPARK_OPEN_PS_NO_LM_HEAD_INBOX": PREFIX + "LM_HEAD_INBOX=off", + "MUSESPARK_OPEN_PS_LAYER0_INPUT": "nothing: model_boundary builds layer 0's attention input", + "MUSESPARK_OPEN_PS_STEP_COMMIT_PDL": "nothing: the step commit launch has no PDL variant", + "MUSESPARK_OPEN_HEAD_INBOX_POLL_SCOPE": "nothing: the inbox polls are system-scope", + "MUSESPARK_OPEN_HEAD_INBOX_RESET_SCOPE": "nothing: the inbox resets are system-scope", + "TRTLLM_MUSESPARK_SWA_BLOCKS_PER_PROGRAM": "nothing: this build has no such switch", + "TRTLLM_MUSESPARK_TEP_WIDE_FULL_FOLD": "nothing: this build has no such switch", + "TRTLLM_MUSESPARK_TEP_WIDE_QKV_PREFETCH": "nothing: this build has no such switch", + "TRTLLM_MUSESPARK_TEP_WIDE_CAPTURE_FOLD": "nothing: this build has no such switch", + } +) +#: Name prefixes only retired switches carry; any other name under them raises too. +RETIRED_PREFIXES: Mapping[str, str] = types.MappingProxyType( + { + "MUSESPARK_OPEN_PS_": "the " + PREFIX + "* knobs", + "TRTLLM_MUSESPARK_TEP_WIDE_": "the " + PREFIX + "WIDE_* knobs", + "MUSESPARK_OPEN_HEAD_INBOX_": "nothing: the inbox scopes are system-scope", + } +) + + +def retired_replacement(name: str) -> Optional[str]: + """What replaced the retired switch ``name``, or None for a name that is not retired.""" + if name in RETIRED: + return RETIRED[name] + for prefix, replacement in RETIRED_PREFIXES.items(): + if name.startswith(prefix): + return replacement + return None + + +class Kind(str, enum.Enum): + """What a lever changes, which fixes what evidence ``auto`` needs (module docstring).""" + + #: A fast path validated with its composition: ``auto`` serves it where the rule says. + COMPOSITION = "composition" + #: Writes the bytes of the path it replaces. + BITWISE = "lossless-bitwise" + #: Same precision, another rounding or summation order: ``auto`` also needs a gate. + NUMERIC = "lossless-numeric" + #: Lower precision: never on under ``auto``. + LOSSY = "lossy" + #: Lossless, with an open safety check: never on under ``auto``. + FLAGGED = "flagged" + + +class Role(str, enum.Enum): + """The server's role: aggregated, disaggregated context or generation, or not known.""" + + AGG = "agg" + CTX = "ctx" + GEN = "gen" + UNKNOWN = "unknown" + + +class Arena(str, enum.Enum): + """Where the DFlash worker keeps the drafter's context K/V.""" + + #: A worker-owned arena of every context column. + DENSE = "dense" + #: A worker-owned arena of each slot's last window (``TRTLLM_DFLASH_WINDOWED_CTX_ARENA``). + WINDOWED = "windowed" + #: Pages of a KV cache manager. + PAGED = "paged" + + +@dataclasses.dataclass(frozen=True) +class Facts: + """The shared-path variables the resolver reads; it never sets or retires them.""" + + #: ``TRTLLM_KV_DEFER_CLOSE`` is exactly ``1``. + kv_defer_close: bool = False + #: ``UCX_RNDV_THRESH`` is set to a non-empty value. + ucx_rndv_thresh_set: bool = False + + def describe(self) -> str: + return ( + f"deferred close {'on' if self.kv_defer_close else 'off'}, " + f"{UCX_RNDV_THRESH_ENV} {'set' if self.ucx_rndv_thresh_set else 'unset'}" + ) + + +@dataclasses.dataclass(frozen=True) +class Deployment: + """The facts a lever's deployment rule reads, fixed when the model is built.""" + + sm: int + tp: int + max_batch_size: Optional[int] + max_num_tokens: Optional[int] + role: Role + #: Where the role came from: ``hint``, ``no cache transceiver``, ``no hint on a + #: disaggregated server`` or ``no engine facts``. + role_source: str + arena: Arena + facts: Facts = Facts() + rows_per_request: int = ROWS_PER_REQUEST + + def decodes(self) -> bool: + """Whether this server runs decode steps (an unknown role counts as one that does).""" + return self.role is not Role.CTX + + def reaches(self, rows: int) -> bool: + """Whether some decode step can have at least ``rows`` rows (true for an unknown + ``max_batch_size``).""" + return self.max_batch_size is None or self.max_batch_size * self.rows_per_request >= rows + + def describe(self) -> str: + return ( + f"sm{self.sm} tep{self.tp}, {self.role.value} ({self.role_source}), max_batch_size " + f"{_count(self.max_batch_size)}, max_num_tokens {_count(self.max_num_tokens)}, " + f"{self.arena.value} arena" + ) + + def as_dict(self) -> dict: + return { + "sm": self.sm, + "tp": self.tp, + "max_batch_size": self.max_batch_size, + "max_num_tokens": self.max_num_tokens, + "role": self.role.value, + "role_source": self.role_source, + "arena": self.arena.value, + } + + +def _count(value: Optional[int]) -> str: + return "unknown" if value is None else str(value) + + +@dataclasses.dataclass(frozen=True) +class Param: + """A heuristic threshold a lever declares, overridable by ``__``. + + An integer in ``low..high``, or, with ``choices``, a non-empty subset of ``choices`` + (comma-separated, kept in ``choices`` order). A lossy lever's range is a subset of + the values its accuracy gate covered. + """ + + key: str + default: object + low: int = 1 + high: int = 1 << 30 + choices: Tuple[str, ...] = () + + def parse(self, raw: str) -> object: + """The value ``raw`` gives this parameter. + + Raises: + ValueError: for a value outside the parameter's range or choices. + """ + text = raw.strip().lower() + if self.choices: + names = [name.strip() for name in text.split(",") if name.strip()] + unknown = sorted(set(names) - set(self.choices)) + if not names or unknown: + raise ValueError( + f"expected a non-empty comma-separated subset of {list(self.choices)}, got " + f"{raw!r}" + ) + return tuple(choice for choice in self.choices if choice in names) + try: + value = int(text, 10) + except ValueError: + raise ValueError( + f"expected an integer in {self.low}..{self.high}, got {raw!r}" + ) from None + if not self.low <= value <= self.high: + raise ValueError(f"expected an integer in {self.low}..{self.high}, got {value}") + return value + + +@dataclasses.dataclass(frozen=True) +class Gate: + """A deployment class where a lossless-numeric lever passed, or must pass, an accuracy gate. + + ``scope`` says the class in words; the evidence lives with the measurements, not here. + """ + + holds: Callable[["Deployment"], bool] + scope: str + + +#: ``(on, the rule in words)`` from a deployment, the decisions so far and the lever's parameters. +RuleFn = Callable[["Deployment", "Resolving", Mapping[str, object]], Tuple[bool, str]] + + +@dataclasses.dataclass(frozen=True) +class Rule: + """How one lever resolves when its rule is not the default (served, on). + + ``origin`` is ``upstream`` for the upstream target's own paths and their upstream + behaviour, else the serving region the lever targets (``HC`` high concurrency, + ``midc`` mid concurrency, ``LL`` low latency): under ``__DEFAULT=baseline`` a + non-upstream lever is off unless its knob is set, and an upstream one runs + ``baseline`` (or ``auto`` when it has none). ``after`` names the decisions ``auto`` / + ``baseline`` read; ``requires`` levers that must be on; + ``conflicts`` levers that may not be on with it; ``excluded`` a structural reason it + cannot be on, from the decisions so far. + """ + + origin: str + kind: Kind + auto: RuleFn + baseline: Optional[RuleFn] = None + params: Tuple[Param, ...] = () + gates: Tuple[Gate, ...] = () + after: Tuple[str, ...] = () + requires: Tuple[str, ...] = () + conflicts: Tuple[str, ...] = () + excluded: Optional[Callable[["Deployment", "Resolving"], Optional[str]]] = None + + def param(self, key: str) -> Param: + for param in self.params: + if param.key == key: + return param + raise KeyError(key) + + +@dataclasses.dataclass(frozen=True) +class Decision: + """One lever's resolution: on or off, by which step, the rule in words, its parameters.""" + + name: str + on: bool + source: str + why: str + params: Mapping[str, object] = dataclasses.field(default_factory=dict) + + +class Resolving: + """The decisions made so far, as a rule reads them (``Rule.after``).""" + + def __init__(self, served: FrozenSet[str], decisions: Mapping[str, Decision]) -> None: + self._served = served + self._decisions = decisions + + def on(self, name: str) -> bool: + """Whether ``name`` resolved on; False for a lever the composition does not serve. + + Raises: + RuntimeError: for a served lever not decided yet (a rule that reads it must + name it in ``Rule.after``). + """ + if name not in self._served: + return False + decision = self._decisions.get(name) + if decision is None: + raise RuntimeError(f"lever rule reads {name}, which is not decided yet") + return decision.on + + def param(self, name: str, key: str) -> object: + """``name``'s parameter ``key`` as resolved, or its default when ``name`` is not served.""" + decision = self._decisions.get(name) + if decision is not None and key in decision.params: + return decision.params[key] + return RULES[name].param(key).default + + +# -------------------------------------------------------------------------- +# The deployment rules +# -------------------------------------------------------------------------- +_NO_DECODE = "the server runs no decode step" + + +def _reaching(rows: int, what: str) -> RuleFn: + """A rule on while the server decodes and some decode step reaches ``rows`` rows.""" + + def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + if not d.decodes(): + return False, _NO_DECODE + if d.reaches(rows): + return True, f"{what} are reachable at max_batch_size {_count(d.max_batch_size)}" + return False, f"max_batch_size {_count(d.max_batch_size)} never reaches {rows} decode rows" + + return rule + + +def _packed_prep_append(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + rows = int(p["MIN_ROWS"]) + return _reaching(rows, f"prep-append calls of {rows} or more rows")(d, t, p) + + +def _windowed_commit_parallel( + d: Deployment, t: Resolving, p: Mapping[str, object] +) -> Tuple[bool, str]: + requests = int(p["MIN_REQUESTS"]) + return _reaching( + d.rows_per_request * requests, f"context commits of {requests} or more requests" + )(d, t, p) + + +def _wide_steps(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + return _reaching( + WIDE_MIN_ROWS, f"decode steps of {WIDE_MIN_ROWS}..{4 * ROWS_PER_REQUEST} rows" + )(d, t, p) + + +def _at_most(key: str, requires_decode: bool = True) -> RuleFn: + """A rule on while ``max_batch_size`` is known and at most the parameter ``key``.""" + + def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + most = int(p[key]) + if requires_decode and not d.decodes(): + return False, _NO_DECODE + if d.max_batch_size is None: + return False, "max_batch_size is unknown" + if d.max_batch_size <= most: + return True, f"max_batch_size {d.max_batch_size} is at most {most}" + return False, f"max_batch_size {d.max_batch_size} is above {most}" + + return rule + + +def _lossy(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + return False, "lossy: on only by its knob" + + +def _flagged(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + return False, "flagged: on only by its knob, once its open safety check has passed" + + +def _always(why: str) -> RuleFn: + def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + return True, why + + return rule + + +def _not_on_sm107(why_off: str, why_on: str) -> RuleFn: + def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + if d.sm == 107: + return False, why_off + return True, why_on + + return rule + + +def _coop_fold(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + if not d.decodes(): + return False, _NO_DECODE + return True, f"head calls of at most {p['MAX_ROWS']} rows" + + +def _hot_replicas(requires_decode: bool) -> RuleFn: + def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + if requires_decode and not d.decodes(): + return False, _NO_DECODE + if t.on("half_window_experts"): + return False, "yields to half_window_experts, which places each step's experts" + return True, "33..512-row decode steps" + + return rule + + +def _attn_gather_excluded(d: Deployment, t: Resolving) -> Optional[str]: + if "qkv" not in t.param("context_fp8_dense", "PROJECTIONS"): + return "its gathered rows feed the QKV projection, which context_fp8_dense__PROJECTIONS leaves out" + return None + + +def _router_gate(role: Role, scope: str) -> Gate: + def holds(d: Deployment) -> bool: + return ( + d.role is role + and d.tp == 4 + and d.max_batch_size is not None + and 2 <= d.max_batch_size <= 4 + ) + + return Gate(holds, scope) + + +#: Every lever whose rule is not the default; every other served lever is on under ``auto`` +#: and ``baseline``. +RULES: Mapping[str, Rule] = types.MappingProxyType( + { + "packed_prep_append": Rule( + "HC", + Kind.BITWISE, + _packed_prep_append, + params=(Param("MIN_ROWS", 33, 8, 1024),), + ), + "windowed_commit_parallel": Rule( + "HC", + Kind.BITWISE, + _windowed_commit_parallel, + params=(Param("MIN_REQUESTS", 5, 1, 64),), + ), + "context_fp8_dense": Rule( + "HC", + Kind.LOSSY, + _lossy, + params=( + Param("MIN_ROWS", 1024, 1024, 32768), + Param( + "PROJECTIONS", + ("qkv", "o_proj", "pool", "unpool"), + choices=("qkv", "o_proj", "pool", "unpool"), + ), + ), + ), + "context_fp8_attn_gather": Rule( + "HC", Kind.LOSSY, _lossy, after=("context_fp8_dense",), excluded=_attn_gather_excluded + ), + "wide_direct_transports": Rule("midc", Kind.BITWISE, _wide_steps), + "wide_router_per_row": Rule( + "midc", + Kind.NUMERIC, + _wide_steps, + gates=( + _router_gate(Role.GEN, "generation servers at tep4 with max_batch_size 2..4"), + _router_gate(Role.AGG, "aggregated servers at tep4 with max_batch_size 2..4"), + ), + ), + "wide_oproj_publish": Rule("midc", Kind.BITWISE, _wide_steps), + "wide_finalize_transport": Rule( + "midc", Kind.BITWISE, _always("17..32-token finalize calls on a 32-token transport") + ), + "lm_head_inbox": Rule( + "upstream", + Kind.COMPOSITION, + _not_on_sm107( + "the sm107 compositions gather the 8-row head without the inbox", "8-row heads" + ), + baseline=_always("8-row heads"), + ), + "graph_start_qkv_prefetch": Rule( + "upstream", + Kind.COMPOSITION, + _not_on_sm107( + "on sm107 a captured decode graph starts without a fork, and layer 0's QKV " + "weight is not prefetched", + "graph starts of split-K row counts", + ), + baseline=_always("graph starts of split-K row counts"), + ), + "step_loop_on_execution_stream": Rule( + "LL", + Kind.BITWISE, + _at_most("MAX_BATCH"), + params=(Param("MAX_BATCH", 4, 1, 64),), + ), + "incremental_block_offsets": Rule( + "LL", + Kind.BITWISE, + _at_most("MAX_BATCH"), + params=(Param("MAX_BATCH", 8, 1, 64),), + ), + "drafter_kv_side_chain": Rule( + "LL", Kind.BITWISE, _always("decode steps of 1..64 requests over a dense arena") + ), + "drafter_fc_early_fork": Rule("LL", Kind.BITWISE, _always("forked side-chain steps")), + "drafter_preamble_from_commit": Rule( + "LL", Kind.BITWISE, _always("folded steps over a dense arena") + ), + "head_sampler_coop_fold": Rule( + "LL", + Kind.BITWISE, + _coop_fold, + params=(Param("MAX_ROWS", 16, 1, 128),), + ), + "head_multimem_gather": Rule( + "LL", + Kind.FLAGGED, + _flagged, + requires=("step_loop_on_execution_stream",), + conflicts=("lm_head_inbox",), + ), + "half_window_experts": Rule( + "upstream", + Kind.COMPOSITION, + _at_most("MAX_BATCH"), + baseline=_at_most("MAX_BATCH", requires_decode=False), + params=(Param("MAX_BATCH", 8, 1, 64),), + ), + "hot_replicas": Rule( + "upstream", + Kind.COMPOSITION, + _hot_replicas(requires_decode=True), + baseline=_hot_replicas(requires_decode=False), + after=("half_window_experts",), + conflicts=("half_window_experts",), + ), + "prefill_cutedsl_experts": Rule( + "upstream", Kind.COMPOSITION, _always("expert calls of 4096..8192 rows") + ), + } +) + + +def default_param(name: str, key: str) -> object: + """Lever ``name``'s declared default for parameter ``key``. + + Raises: + KeyError: when ``name`` declares no parameter ``key``. + """ + rule = RULES.get(name) + if rule is None: + raise KeyError(f"lever {name} declares no parameter {key}") + return rule.param(key).default + + +def knob_name(name: str, key: Optional[str] = None) -> str: + """The environment variable of lever ``name`` (or of its parameter ``key``).""" + return PREFIX + name.upper() + ("" if key is None else SEP + key) + + +# -------------------------------------------------------------------------- +# The environment +# -------------------------------------------------------------------------- +def snapshot_of(environ: Mapping[str, str]) -> Dict[str, str]: + """The variables of ``environ`` the resolver reads: lever knobs, retired names, facts. + + A lever knob set to an empty value is left out (unset); a retired name is kept + whatever its value. + """ + snapshot: Dict[str, str] = {} + for name, value in environ.items(): + if name.startswith(PREFIX): + if value.strip(): + snapshot[name] = value + elif retired_replacement(name) is not None: + snapshot[name] = value + elif name in (KV_DEFER_CLOSE_ENV, UCX_RNDV_THRESH_ENV): + snapshot[name] = value + return snapshot + + +_SNAPSHOT: Optional[Mapping[str, str]] = None + + +def environment() -> Mapping[str, str]: + """The snapshot of this process's environment, read on the first call and never again.""" + global _SNAPSHOT + if _SNAPSHOT is None: + _SNAPSHOT = types.MappingProxyType(snapshot_of(os.environ)) + return _SNAPSHOT + + +def facts_of(snapshot: Mapping[str, str]) -> Facts: + """The two shared-path facts of ``snapshot``.""" + return Facts( + kv_defer_close=snapshot.get(KV_DEFER_CLOSE_ENV) == "1", + ucx_rndv_thresh_set=bool(snapshot.get(UCX_RNDV_THRESH_ENV, "").strip()), + ) + + +#: The prefix of every item ``tagged`` lists for the TP agreement. +AGREEMENT_TAG = "lever environment: " + + +def tagged(snapshot: Mapping[str, str]) -> List[str]: + """``snapshot`` as items for ``dispatch.agreed_union``: every rank must list the same. + + The knobs and the retired names with their raw values, the facts as resolved. + """ + facts = facts_of(snapshot) + items = [ + f"{AGREEMENT_TAG}{name}={value}" + for name, value in sorted(snapshot.items()) + if name not in (KV_DEFER_CLOSE_ENV, UCX_RNDV_THRESH_ENV) + ] + items.append(f"{AGREEMENT_TAG}{facts.describe()}") + return items + + +@dataclasses.dataclass(frozen=True) +class Knobs: + """The lever knobs of a snapshot, parsed and validated.""" + + #: Lever name -> ``auto`` | ``on`` | ``off``. + modes: Mapping[str, str] + #: (lever name, parameter key) -> parsed value. + params: Mapping[Tuple[str, str], object] + default: str + role: str + #: The knobs as set (names under ``PREFIX``). + raw: Mapping[str, str] + + +def parse_knobs(snapshot: Mapping[str, str], names: Iterable[str]) -> Knobs: + """The lever knobs of ``snapshot`` for a registry of lever ``names``. + + Raises: + ValueError: naming every retired switch that is set, every unknown name under + ``PREFIX``, every value outside its set and every parameter outside its range + or on a lever that declares no such parameter. + """ + names = frozenset(names) + errors: List[str] = [] + for name in sorted(snapshot): + replacement = retired_replacement(name) + if replacement is not None: + errors.append(f"{name} is retired (use {replacement})") + modes: Dict[str, str] = {} + params: Dict[Tuple[str, str], object] = {} + default, role = AUTO, AUTO + raw: Dict[str, str] = {} + for name in sorted(snapshot): + if not name.startswith(PREFIX): + continue + value = snapshot[name].strip().lower() + raw[name] = snapshot[name].strip() + rest = name[len(PREFIX) :] + if rest == "_DEFAULT": + if value not in DEFAULT_VALUES: + errors.append(f"{name}={value!r}: expected one of {list(DEFAULT_VALUES)}") + default = value + continue + if rest == "_ROLE": + allowed = [AUTO] + [r.value for r in Role if r is not Role.UNKNOWN] + if value not in allowed: + errors.append(f"{name}={value!r}: expected one of {allowed}") + role = value + continue + lever_part, sep, key = rest.partition(SEP) + lever = lever_part.lower() + if rest.startswith("_") or lever not in names or lever.upper() != lever_part: + close = difflib.get_close_matches(lever, sorted(names), n=1) + hint = f" (did you mean {knob_name(close[0])}?)" if close else "" + errors.append(f"{name}: no lever of this target is named {lever!r}{hint}") + continue + if not sep: + if value not in LEVER_VALUES: + errors.append(f"{name}={value!r}: expected one of {list(LEVER_VALUES)}") + else: + modes[lever] = value + continue + rule = RULES.get(lever) + declared = [param.key for param in rule.params] if rule is not None else [] + if key not in declared: + errors.append( + f"{name}: lever {lever} declares " + + (f"the parameters {declared}" if declared else "no parameter") + ) + continue + try: + params[(lever, key)] = rule.param(key).parse(snapshot[name]) + except ValueError as error: + errors.append(f"{name}: {error}") + if errors: + raise ValueError("; ".join(errors)) + return Knobs(modes, params, default, role, raw) + + +def no_composition_notes(snapshot: Mapping[str, str], names: Iterable[str]) -> List[str]: + """For a deployment without a composition: validate ``snapshot``, then say what it sets. + + Raises: + ValueError: as ``parse_knobs``. + """ + knobs = parse_knobs(snapshot, names) + return [ + f"{name}={value} has no effect without a composition" for name, value in knobs.raw.items() + ] + + +# -------------------------------------------------------------------------- +# The deployment +# -------------------------------------------------------------------------- +def arena_of(spec_worker, draft_model, paged_backends: Sequence[str]) -> Arena: + """The drafter context arena the DFlash worker will keep, by its own predicate. + + ``paged`` when the drafter attention backend is one of ``paged_backends`` + (``speculative.dflash._PAGED_ATTENTION_BACKENDS``) or the drafter asks for pages + (``_paged_ctx_cache``); else ``windowed`` when the worker windows its arena; else + ``dense``. An ``AUTO`` worker backend is the drafter's, as the worker resolves it at + the drafter's first forward. + """ + backend = getattr(spec_worker, "_dflash_attention_backend", None) + if backend == "AUTO": + backend = getattr(draft_model, "dflash_attention_backend", backend) + if backend in tuple(paged_backends) or getattr(draft_model, "_paged_ctx_cache", False): + return Arena.PAGED + if getattr(spec_worker, "windowed_ctx_arena", False): + return Arena.WINDOWED + return Arena.DENSE + + +def worker_arena(spec_worker) -> Optional[Arena]: + """The arena the DFlash worker built, once it has (None before its first drafter step).""" + if not getattr(spec_worker, "_ctx_buf_inited", False): + return None + if getattr(spec_worker, "_ctx_paged", False): + return Arena.PAGED + if getattr(spec_worker, "_ctx_base", None) is not None: + return Arena.WINDOWED + return Arena.DENSE + + +#: The role source of a server with a cache transceiver and no role hint. +DISAGGREGATED_UNHINTED = "no hint on a disaggregated server" + + +def role_of(snapshot: Mapping[str, str], disaggregated: Optional[bool]) -> Tuple[Role, str]: + """The server's role and where it came from. + + The hint ``ROLE_KNOB`` wins when set. Otherwise ``agg`` without a cache transceiver, + and ``unknown`` with one or without engine facts. + + Raises: + ValueError: for a hint the server contradicts: ``agg`` with a cache transceiver, + ``ctx`` or ``gen`` without one. + """ + hint = snapshot.get(ROLE_KNOB, AUTO).strip().lower() or AUTO + if hint != AUTO: + role = Role(hint) + if role is Role.AGG and disaggregated is True: + raise ValueError(f"{ROLE_KNOB}=agg on a server with a cache transceiver") + if role in (Role.CTX, Role.GEN) and disaggregated is False: + raise ValueError(f"{ROLE_KNOB}={hint} on a server without a cache transceiver") + return role, "hint" + if disaggregated is None: + return Role.UNKNOWN, "no engine facts" + if disaggregated: + return Role.UNKNOWN, DISAGGREGATED_UNHINTED + return Role.AGG, "no cache transceiver" + + +def deployment_of( + model_config, compute_capability, arena: Arena, engine_facts, snapshot: Mapping[str, str] +) -> Deployment: + """The ``Deployment`` of a model: its config's mapping, the device, the arena and the engine. + + ``engine_facts`` carries ``max_batch_size``, ``max_num_tokens`` and ``disaggregated`` + (``_router_index.DeploymentFacts``), or is None when the target class was built without + them. + + Raises: + ValueError: as ``role_of``. + """ + major, minor = compute_capability + role, source = role_of(snapshot, getattr(engine_facts, "disaggregated", None)) + return Deployment( + sm=10 * int(major) + int(minor), + tp=int(model_config.mapping.tp_size), + max_batch_size=getattr(engine_facts, "max_batch_size", None), + max_num_tokens=getattr(engine_facts, "max_num_tokens", None), + role=role, + role_source=source, + arena=arena, + facts=facts_of(snapshot), + ) + + +# -------------------------------------------------------------------------- +# The table +# -------------------------------------------------------------------------- +@dataclasses.dataclass(frozen=True) +class LeverTable: + """Every served lever's decision, the deployment it was resolved for, and the knobs set.""" + + deployment: Deployment + #: Served lever name -> its decision, in registry order. + decisions: Mapping[str, Decision] + #: The rule-table levers the composition does not serve. + unserved: FrozenSet[str] + #: The knobs as set (names under ``PREFIX``). + knobs: Mapping[str, str] + default: str = AUTO + #: Things the log should say as warnings (an unknown role, ``on`` outside a gate, ...). + warnings: Tuple[str, ...] = () + + def on(self, name: str) -> bool: + decision = self.decisions.get(name) + return decision is not None and decision.on + + def param(self, name: str, key: str) -> object: + """``name``'s resolved parameter ``key``, or its default when ``name`` is not served. + + Raises: + KeyError: when ``name`` declares no parameter ``key``. + """ + decision = self.decisions.get(name) + if decision is not None and key in decision.params: + return decision.params[key] + return default_param(name, key) + + def paths(self) -> FrozenSet[str]: + """The levers on: closed under ``depends_on`` and ``Rule.requires`` by construction.""" + return frozenset(name for name, decision in self.decisions.items() if decision.on) + + def describe(self) -> Tuple[str, ...]: + """The log lines of the table (without the log tag).""" + knobs = ", ".join(f"{name}={value}" for name, value in self.knobs.items()) or "none" + lines = [ + f"levers: {self.deployment.describe()}; default {self.default}; facts: " + f"{self.deployment.facts.describe()}; knobs set: {knobs}" + ] + for name, decision in self.decisions.items(): + rule = RULES.get(name) + if rule is None and decision.on: + continue + kind, origin = (rule.kind.value, rule.origin) if rule else ("composition", "upstream") + lines.append( + f"lever {name} {ON if decision.on else OFF} ({decision.source}: {decision.why})" + + (f" {dict(decision.params)}" if decision.params else "") + + f" [{kind}, {origin}]" + ) + if self.unserved: + lines.append( + f"levers not served at sm{self.deployment.sm} tep{self.deployment.tp}: " + + ", ".join(name for name in RULES if name in self.unserved) + ) + lines.append(f"lever table json (resolved): {self.as_json()}") + return tuple(lines) + + def as_json(self, installed: Optional[Mapping[str, bool]] = None) -> str: + """The table as one JSON line; with ``installed``, each lever's install outcome too. + + A lever without a reported outcome counts as installed when it resolved on. + """ + levers = {} + params = {} + for name, decision in self.decisions.items(): + entry = {"resolved": ON if decision.on else OFF, "source": decision.source} + if installed is not None: + entry["installed"] = bool(installed.get(name, decision.on)) + levers[name] = entry + if decision.params: + params[name] = { + key: list(value) if isinstance(value, tuple) else value + for key, value in decision.params.items() + } + table = { + "deployment": self.deployment.as_dict(), + "facts": dataclasses.asdict(self.deployment.facts), + "default": self.default, + "knobs": dict(self.knobs), + "params": params, + "levers": levers, + } + return json.dumps(table, sort_keys=True, separators=(",", ":")) + + +def _order( + served: FrozenSet[str], + depends_on: Mapping[str, Sequence[str]], + requires: Mapping[str, Sequence[str]], +) -> List[str]: + """``served`` in an order where every lever follows what it builds on, requires or reads. + + Ties keep the registry order of ``depends_on``. + """ + registry = [name for name in depends_on if name in served] + edges = { + name: { + other + for other in ( + *depends_on[name], + *requires.get(name, ()), + *(RULES[name].requires if name in RULES else ()), + *(RULES[name].after if name in RULES else ()), + ) + if other in served and other != name + } + for name in registry + } + ordered: List[str] = [] + done: set = set() + while len(ordered) < len(registry): + ready = [name for name in registry if name not in done and edges[name] <= done] + if not ready: + raise RuntimeError( + f"lever dependencies form a cycle among {sorted(set(registry) - done)}" + ) + ordered.append(ready[0]) + done.add(ready[0]) + return ordered + + +def resolve( + deployment: Deployment, + served: Iterable[str], + structural: Mapping[str, str], + snapshot: Mapping[str, str], + depends_on: Mapping[str, Sequence[str]], + requires: Optional[Mapping[str, Sequence[str]]] = None, +) -> LeverTable: + """The lever table of ``deployment`` for the composition's ``served`` levers. + + ``structural`` maps each lever a structural rule excludes to the reason; + ``depends_on`` is the registry (``dispatch.DEPENDS_ON``: every lever and what it + builds on); ``requires`` adds the deployment's own requirements + (``dispatch.structural_rules``). See the module docstring for the order. + + Raises: + ValueError: on every error of the module docstring, all named in one message. + """ + knobs = parse_knobs(snapshot, depends_on) + requires = dict(requires or {}) + served = frozenset(served) + d = deployment + errors: List[str] = [] + warnings: List[str] = [] + transfers_kv = d.role in (Role.GEN, Role.CTX) or ( + d.role is Role.UNKNOWN and d.role_source == DISAGGREGATED_UNHINTED + ) + if d.facts.ucx_rndv_thresh_set and transfers_kv: + errors.append( + f"{UCX_RNDV_THRESH_ENV} is set on a {d.role.value} server with a cache transceiver: " + "it would move this process's KV-transfer rendezvous too" + ) + if d.role is Role.UNKNOWN: + warnings.append( + f"role unknown ({d.role_source}): counted as a decoding server; set {ROLE_KNOB}" + ) + decisions: Dict[str, Decision] = {} + t = Resolving(served, decisions) + for name in _order(served, depends_on, requires): + rule = RULES.get(name) + mode = knobs.modes.get(name) + params = { + param.key: knobs.params.get((name, param.key), param.default) + for param in (rule.params if rule is not None else ()) + } + decision = _decide( + name, rule, mode, knobs.default, params, d, t, structural, depends_on, requires + ) + if mode == ON and not decision.on: + errors.append(f"{knob_name(name)}=on, but {decision.why}") + if ( + mode == ON + and decision.on + and rule is not None + and rule.kind is Kind.NUMERIC + and not any(gate.holds(d) for gate in rule.gates) + ): + warnings.append( + f"{knob_name(name)}=on outside its accuracy gates (" + + "; ".join(gate.scope for gate in rule.gates) + + ")" + ) + decisions[name] = decision + for name, decision in decisions.items(): + rule = RULES.get(name) + if not decision.on or rule is None: + continue + for other in rule.conflicts: + if other in decisions and decisions[other].on: + errors.append( + f"{name} ({decision.source}) and {other} ({decisions[other].source}) cannot " + f"both be on; set {knob_name(other)}=off or {knob_name(name)}=off" + ) + # The MXFP8 copy takes any call of MIN_ROWS rows or more, context rows or not. + fp8 = decisions.get("context_fp8_dense") + if fp8 is not None and fp8.on and d.decodes() and d.reaches(int(fp8.params["MIN_ROWS"])): + errors.append( + f"{knob_name('context_fp8_dense')}=on, but decode steps of {fp8.params['MIN_ROWS']} " + f"rows or more are reachable at max_batch_size {_count(d.max_batch_size)}, and the " + f"MXFP8 copy would take them too; set {knob_name('context_fp8_dense', 'MIN_ROWS')} " + f"above max_batch_size * {d.rows_per_request}" + ) + for name, mode in knobs.modes.items(): + if name in served: + continue + if mode == ON: + errors.append( + f"{knob_name(name)}=on, but the composition at sm{d.sm} tep{d.tp} does not " + f"serve {name}" + ) + else: + warnings.append(f"{knob_name(name)}={mode} has no effect: {name} is not served here") + for name, key in knobs.params: + if name not in served: + warnings.append(f"{knob_name(name, key)} has no effect: {name} is not served here") + if errors: + raise ValueError("; ".join(errors)) + ordered = {name: decisions[name] for name in depends_on if name in decisions} + return LeverTable( + deployment=d, + decisions=types.MappingProxyType(ordered), + unserved=frozenset(RULES) - served, + knobs=types.MappingProxyType(dict(knobs.raw)), + default=knobs.default, + warnings=tuple(warnings), + ) + + +def _decide( + name: str, + rule: Optional[Rule], + mode: Optional[str], + default: str, + params: Mapping[str, object], + d: Deployment, + t: Resolving, + structural: Mapping[str, str], + depends_on: Mapping[str, Sequence[str]], + requires: Mapping[str, Sequence[str]], +) -> Decision: + """One lever's decision, from the decisions already made (``resolve``'s order).""" + if name in structural: + return Decision(name, False, "structural", structural[name], params) + for other in depends_on[name]: + if not t.on(other): + return Decision( + name, False, "dependency", f"it builds on {other}, which is off", params + ) + for other in (*requires.get(name, ()), *(rule.requires if rule is not None else ())): + if not t.on(other): + return Decision(name, False, "dependency", f"it requires {other}, which is off", params) + if rule is not None and rule.excluded is not None: + reason = rule.excluded(d, t) + if reason is not None: + return Decision(name, False, "structural", reason, params) + if mode == OFF: + return Decision(name, False, "knob", f"{knob_name(name)}=off", params) + if mode == ON: + return Decision(name, True, "knob", f"{knob_name(name)}=on", params) + baseline = mode is None and default == BASELINE + if rule is None: + return Decision(name, True, BASELINE if baseline else AUTO, "served", params) + if baseline: + if rule.baseline is not None: + on, why = rule.baseline(d, t, params) + elif rule.origin == "upstream": + on, why = rule.auto(d, t, params) + else: + on, why = False, f"not an upstream lever ({rule.origin}): off unless its knob is set" + return Decision(name, on, BASELINE, why, params) + on, why = rule.auto(d, t, params) + if on and rule.kind is Kind.NUMERIC: + held = [gate for gate in rule.gates if gate.holds(d)] + if not held: + scopes = "; ".join(gate.scope for gate in rule.gates) + return Decision( + name, False, "gate", f"{why}, but no accuracy gate holds ({scopes})", params + ) + why = f"{why}; gate: {held[0].scope}" + return Decision(name, on, AUTO, why, params) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py index f99ce5c8bae3..3d802268755e 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/modeling.py @@ -190,7 +190,11 @@ class acts only on sliding-window layers): their decode prep also writes the Other row counts keep the hop (``DirectPublish``). At those 8 rows o_proj forks no prefetch of the pool weight. After the weights load, ``verify_direct_publish`` checks the kernel the published o_proj captures; a - step that raises while armed resets the transport. + step that raises while armed resets the transport. With ``wide_post_attention_hop`` + and the ``wide_oproj_publish`` lever, the + TP4 decode steps of 16, 24 and 32 rows publish the same way into a 32-row direct + transport (``layers.WIDE_PUBLISH_ROWS``, the ``wide_dense_gemms`` o_proj tactics) instead + of running the wide hop; same bits. * ``multi_request_post_attention_hop`` (sm103 tep4 only) arms the same joins on decode steps of 9..1024 rows (two to 128 requests at K = 7) with ``ops/allreduce/post_attn_hop_wide``: the o_proj partial skips the built-in MNNVL all-reduce and one kernel reduces it, joins @@ -257,6 +261,13 @@ class acts only on sliding-window layers): their decode prep also writes the of 1..16 tokens run ``ops/moe/finalize_allreduce_transport``, bitwise the built-in kernel's result on a different transport; larger calls run the built-in kernel. No module is re-classed. +* ``wide_finalize_transport`` (tep4 on sm_107 only) has that workspace build a + second one on a 32-token layout (its own receive area, arrival counter and + kernel build): calls of 17..32 tokens -- decode steps of three or four + requests at K = 7 -- run the same kernel there, bitwise the built-in TP4 + kernel's result (one CTA per token at every count up to its 128), and only + calls above 32 tokens run the built-in kernel. Calls of 1..16 tokens keep the + 16-token workspace. * ``router_gemm`` re-classes every MoE layer once more, with ``RouterGemmMoE`` in front of the class the paths above left it with, and ``DecodeTarget`` arms it on decode steps of 9..512 rows (two or more requests, which the 8-row prologue does @@ -320,12 +331,12 @@ class acts only on sliding-window layers): their decode prep also writes the * ``three_window_placement`` stacks a third ring window on them. * ``half_window_experts`` (tp 4) stacks the ring predecessor's upper expert half on them, which the prologue's routed kernel places onto (and, on calls of 9..32 rows - that reach the routing method, ``ring_buddy_assign``'s half flow); only up to the - ``max_batch_size`` of its capacity gate (``dispatch.CAPACITY_GATES``). + that reach the routing method, ``ring_buddy_assign``'s half flow); only up to its + lever's ``MAX_BATCH`` (``levers.RULES``). * ``bf16_routed_experts`` gives the two BF16 layers' stacked calls unpacked routing (on the owner-only window as well). -* ``hot_replicas`` (tep4, in a plan without ``half_window_experts``: - ``dispatch.YIELDS``) hooks the NVFP4 layers' expert and routing +* ``hot_replicas`` (tep4, in a plan without ``half_window_experts``: the two + levers conflict) hooks the NVFP4 layers' expert and routing calls after the owner-window install; the copies themselves are gathered on the first forward of a deployment whose largest decode step (``max_batch_size`` times the verify width) reaches 33 rows (``HotReplicas.decide``), and ``DecodeTarget`` @@ -409,6 +420,7 @@ class acts only on sliding-window layers): their decode prep also writes the import torch from tensorrt_llm._torch._experimental.modeling_v2._router_index import ( + resolved_facts, resolved_max_batch_size, resolved_mode, ) @@ -418,6 +430,7 @@ class acts only on sliding-window layers): their decode prep also writes the from tensorrt_llm._torch.models.modeling_musespark_open_omni import ( MuseSparkOpenOmniForConditionalGeneration, ) +from tensorrt_llm.logger import logger from . import dispatch as _dispatch from . import layers as _layers @@ -532,6 +545,10 @@ def forward( # The first forward decides whether this deployment allocates the hot # expert copies (kernels/replicated_experts/hot_replicas.py). hot.decide(attn_metadata, kwargs.get("spec_metadata")) + check = self.__dict__.get("_arena_check") + if check is not None and check(): + # The worker built its drafter arena, and it is the one the plan resolved for. + self.__dict__["_arena_check"] = None return self._select_target(attn_metadata).forward( attn_metadata, input_ids=input_ids, @@ -621,6 +638,22 @@ def forward( return hidden_states +def _early_fc_index(fork, spec_metadata, num_layers: int) -> int: + """The decoder-layer index at whose start ``fork.fire_early()`` runs, or -1. + + The index after the last DFlash capture layer (``spec_metadata.layers_to_capture``), when + the fork carries ``drafter_fc_early_fork`` (``ContextKvFork.early_fc_on``) and that index + is a decoder layer (all capture rows are final when it starts). + """ + if fork is None or not fork.early_fc_on or spec_metadata is None: + return -1 + capture = getattr(spec_metadata, "layers_to_capture", None) + if not capture or min(capture) < 0: + return -1 + index = max(capture) + 1 + return index if 0 < index < num_layers else -1 + + class DecodeTarget: """The specialization: routed to only when there are no context rows. @@ -669,11 +702,12 @@ class DecodeTarget: With ``pool_into_route``, a forward of exactly 8 rows captured into a CUDA graph arms every routed layer's pool transport (``DirectPool.armed``): the pool GEMM writes into the layer's pool slot, the forked router shard into its - logit boxes, and the routed pool kernel only polls them. + logit boxes, and the routed pool kernel only polls them. With ``wide_prologue`` + a captured forward of 16..32 rows arms the wide transport's sites the same way. - With ``next_qkv_prefetch``, the forward first forks the L2 prefetch of layer - 0's QKV weight (``QkvPrefetch.graph_start``); every other layer's is forked by - the previous layer's expert call. + With ``next_qkv_prefetch`` every other layer's QKV weight L2 prefetch is forked by + the previous layer's expert call; with ``graph_start_qkv_prefetch`` the forward first + forks layer 0's (``QkvPrefetch.graph_start``). With ``router_gemm``, a step of 9..512 rows arms the router GEMM (``RouterGemmCall.armed``): every MoE layer's router logits run @@ -727,6 +761,7 @@ def __init__(self, core: MuseSparkOpenTepModelingV2) -> None: self.unpool_direct = getattr(core, "_unpool_direct", None) self.pool_direct = getattr(core, "_pool_direct", None) self.qkv_prefetch = getattr(core, "_qkv_prefetch", None) + self.graph_start_prefetch = core._plan.enabled("graph_start_qkv_prefetch") self.router_gemm = getattr(core, "_router_gemm", None) self.kv_splits = getattr(core, "_decode_kv_splits", None) self.wide_moe_allreduce = getattr(core, "_wide_moe_allreduce", None) @@ -747,7 +782,7 @@ def forward( core = self.core if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("Specify exactly one of input_ids and inputs_embeds") - if self.qkv_prefetch is not None: + if self.qkv_prefetch is not None and self.graph_start_prefetch: # Layer 0's QKV weight, asked of the L2 while the embedding chain runs. self.qkv_prefetch.graph_start( int((input_ids if input_ids is not None else inputs_embeds).shape[0]) @@ -786,6 +821,8 @@ def forward( fork = getattr(core, "_context_kv_fork", None) if fork is not None: fork.arm(kwargs.get("spec_metadata"), attn_metadata) + # drafter_fc_early_fork: the layer before which the side chain's local part forks. + early_index = _early_fc_index(fork, kwargs.get("spec_metadata"), len(core.layers)) direct = self.unpool_direct rows = hidden_states.shape[0] @@ -833,6 +870,8 @@ def forward( model_boundary.armed() if model_boundary is not None else contextlib.nullcontext(), ): for index, decoder_layer in enumerate(layers): + if index == early_index: + fork.fire_early() next_norm = layers[index + 1].pre_qkv_proj_norm if index + 1 < len(layers) else None result = decoder_layer( position_ids=position_ids, @@ -912,11 +951,15 @@ def forward( ) #: The fast paths validated at a tep4 rank's shape on sm107 only: the generation rows of #: mixed steps on the multi-request decode FMHA (on sm103 flashinfer's trtllm-gen FMHA -#: already serves both phases of a mixed step). -_SM107_TEP4_ONLY_PATHS = frozenset({"mixed_step_gen_rows"}) -#: TEP8: every fast path but the TP4-only ones of either architecture. +#: already serves both phases of a mixed step), and the finalize transport's 32-token +#: layout for the expert finalize + all-reduce of 17..32-token calls (decode steps of +#: three or four requests at K = 7, which the 16-token transport hands to the built-in +#: kernel). +_SM107_TEP4_ONLY_PATHS = frozenset({"mixed_step_gen_rows", "wide_finalize_transport"}) +#: TEP8: every base fast path (``dispatch.BASE_PATH_NAMES``) but the TP4-only ones of +#: either architecture. _TEP8_PATHS = ( - _dispatch.FAST_PATH_NAMES - _TEP4_ONLY_PATHS - _SM103_TEP4_ONLY_PATHS - _SM107_TEP4_ONLY_PATHS + _dispatch.BASE_PATH_NAMES - _TEP4_ONLY_PATHS - _SM103_TEP4_ONLY_PATHS - _SM107_TEP4_ONLY_PATHS ) #: The fast paths written for any rank count whose output is bitwise the built-in #: model's. @@ -959,9 +1002,9 @@ def forward( #: prologue's routed kernel, those of 9..32 rows (two to four requests, behind the #: multi-request router GEMM and projection gathers) by ``ring_buddy_assign``'s half #: flow, both balancing the selected experts over the four ranks. As on sm107 the -#: window costs 45 GiB of KV cache per rank and is carried only up to -#: ``max_batch_size`` 8 (``dispatch.CAPACITY_GATES``); a plan carrying it leaves -#: ``hot_replicas`` out (``dispatch.YIELDS``). +#: window costs 45 GiB of KV cache per rank and is carried only up to its lever's +#: ``MAX_BATCH`` (``levers.RULES``); a plan carrying it leaves ``hot_replicas`` out (the +#: two levers conflict). _SM103_TEP4_PATHS = ( _TEP4_PATHS | _SM103_TEP4_ONLY_PATHS @@ -1020,7 +1063,7 @@ def forward( #: hold their ring predecessor's upper expert half (``half_window_experts``), so the #: routed kernel balances the selected experts of a step of at most 8 rows over the #: four ranks; that window costs 45 GiB of KV cache per rank, so a plan carries it -#: only up to ``max_batch_size`` 8 (``dispatch.CAPACITY_GATES``). The sharded drafter fc +#: only up to its lever's ``MAX_BATCH`` (``levers.RULES``). The sharded drafter fc #: of up to 128 rows is gathered over the tray's NVLink (``drafter_fc_inbox``). Decode #: steps of two to four #: requests (16..32 rows) also take the TP4 sharded MoE prologue and projection gathers @@ -1076,8 +1119,8 @@ def forward( #: rows stay on ``wide_prologue``'s router shard), with one CTA per SM of the device; #: the attention block's QKV TGV GEMM on the rows of the R200 table and the PDL #: norm-gates of 33..512-row decode steps (``batch_dense_gemms``); and the hot NVFP4 -#: expert copies of 33..512-row decode steps, which yield to ``half_window_experts`` -#: wherever its capacity gate keeps it (``dispatch.YIELDS``). +#: expert copies of 33..512-row decode steps, which conflict with ``half_window_experts`` +#: and stay out of a plan that carries it (``levers.RULES``). _SM107_TEP4_GEMM_EXPERT_PATHS = frozenset({"router_gemm", "batch_dense_gemms", "hot_replicas"}) _SM107_TEP4_PATHS = _SM107_TEP4_PATHS | _SM107_TEP4_GEMM_EXPERT_PATHS #: The sm103 tep4 decode-prep and drafter-at-batch paths sm107 TEP4 serves too, above @@ -1104,14 +1147,62 @@ def forward( #: other sm107 TEP4 paths. _SM107_TEP4_CONTEXT_PATHS = frozenset({"context_sequence_parallel", "context_prep"}) +#: The levers registered beside the fast paths (``dispatch.LEVER_NAMES``) that a +#: composition serving ``next_qkv_prefetch`` serves with it: layer 0's QKV weight prefetch +#: forked at the start of a decode forward. +_PREFETCH_LEVERS = frozenset({"graph_start_qkv_prefetch"}) +#: The levers both sm107 compositions serve on top of their paths: the executor's step loop +#: on the execution stream and its incremental block-offset copy, the windowed commit's +#: parallel form (a windowed arena) and the head sampler's cooperative kernel A. Each +#: lever's rule (``levers.RULES``) decides where it is on; sm103 serves no lever beyond +#: the upstream target's. +_SM107_LEVERS = frozenset( + { + "step_loop_on_execution_stream", + "incremental_block_offsets", + "windowed_commit_parallel", + "head_sampler_coop_fold", + } +) +#: The levers sm107 TEP8 adds: the drafter's fused context-K/V side chain, its early fc fork +#: and the drafter preamble written by the folded commit (a dense arena), and the 8-row head +#: logits gathered over MULTIMEM. +_SM107_TEP8_LEVERS = frozenset( + { + "drafter_kv_side_chain", + "drafter_fc_early_fork", + "drafter_preamble_from_commit", + "head_multimem_gather", + } +) +#: The levers sm107 TEP4 adds, on its TP4 paths: the multi-request prep append's packed +#: kernel, the 16..32-row steps' direct transports, per-row router shard and o_proj +#: publication, the lossy MXFP8 dense projections of large context calls and the MXFP8 +#: gather of the sequence-parallel attention input, and the NVFP4 experts of large prefill +#: calls on the Rubin CuTe DSL GEMMs. +_SM107_TEP4_LEVERS = frozenset( + { + "packed_prep_append", + "wide_direct_transports", + "wide_router_per_row", + "wide_oproj_publish", + "context_fp8_dense", + "context_fp8_attn_gather", + "prefill_cutedsl_experts", + } +) + #: The compositions, per (SM, tensor-parallel size) (``dispatch.CompositionKey``): -#: the fast paths served and the GPU topology they were validated on, with DFlash -#: drafting a linear chain of ``dispatch.DRAFT_LEN`` tokens on an NVFP4 -#: checkpoint; every other deployment runs the built-in modules only. +#: the fast paths and levers served and the GPU topology they were validated on, with +#: DFlash drafting a linear chain of ``dispatch.DRAFT_LEN`` tokens on an NVFP4 +#: checkpoint; every other deployment runs the built-in modules only. The lever table +#: (``dispatch.compose``) decides which served levers a deployment is built with. TARGET_PATHS = { - (103, 8): _dispatch.TargetPaths(served=_TEP8_PATHS, topology=_dispatch.TWO_TRAYS_OF_4), + (103, 8): _dispatch.TargetPaths( + served=_TEP8_PATHS | _PREFETCH_LEVERS, topology=_dispatch.TWO_TRAYS_OF_4 + ), (107, 8): _dispatch.TargetPaths( - served=_TEP8_PATHS, + served=_TEP8_PATHS | _PREFETCH_LEVERS | _SM107_LEVERS | _SM107_TEP8_LEVERS, topology=_dispatch.TWO_TRAYS_OF_4, preconditions=(_dispatch.SM107_FUSED_KERNELS,), ), @@ -1121,13 +1212,49 @@ def forward( preconditions=(_dispatch.NO_PREFILL_SEQUENCE_PARALLEL,), ), (107, 4): _dispatch.TargetPaths( - served=_SM107_TEP4_PATHS | _SM107_TEP4_CONTEXT_PATHS, + served=( + _SM107_TEP4_PATHS + | _SM107_TEP4_CONTEXT_PATHS + | _PREFETCH_LEVERS + | _SM107_LEVERS + | _SM107_TEP4_LEVERS + ), topology=_dispatch.ONE_TRAY_OF_4, preconditions=(_dispatch.NO_PREFILL_SEQUENCE_PARALLEL, _dispatch.SM107_FUSED_KERNELS), ), } +def _arena_check(spec_worker, planned): + """A check of the drafter context arena the DFlash worker builds against ``planned``. + + The plan's arena-bound levers were resolved for ``planned`` (``levers.arena_of``, + ``dispatch.structural_rules``); the worker builds its arena at its first drafter step. + The returned callable answers False until the worker has built it + (``levers.worker_arena``), then True when it is ``planned``, and raises otherwise. + None without a worker. + + Raises (from the callable): + RuntimeError: when the worker built another arena than the plan's. + """ + if spec_worker is None: + return None + from . import levers as _levers + + def check() -> bool: + built = _levers.worker_arena(spec_worker) + if built is None: + return False + if built is not planned: + raise RuntimeError( + f"{_dispatch.TAG} the DFlash worker built a {built.value} drafter context arena; " + f"the plan's levers were resolved for a {planned.value} one" + ) + return True + + return check + + def device_compute_capability() -> Optional[Tuple[int, int]]: """This rank's compute capability, the fact ``route`` routes on; None without a CUDA device. @@ -1162,11 +1289,15 @@ def __init__(self, model_config: ModelConfig, *args, **kwargs): before the built-in constructor runs. The preconditions read the built model (and the drafter arena its DFlash worker keeps), so the plan is decided right after it, before anything is installed - (``dispatch.compose``, agreed over the TP group). + (``dispatch.compose``: the lever table over the engine facts + ``_router_index.resolved_facts`` recorded, agreed over the TP group). The + drafter arena the worker builds at its first drafter step is checked + against the plan's at the first forward after it. Raises: ValueError: under ``require`` without a composition; under any mode - with a precondition of the plan missing (``dispatch.compose``). + with a lever error or a precondition of the plan missing + (``dispatch.compose``). """ mode = resolved_mode(type(self)) capability = device_compute_capability() @@ -1180,6 +1311,8 @@ def __init__(self, model_config: ModelConfig, *args, **kwargs): capability, spec_worker=getattr(self.llm, "spec_worker", None), max_batch_size=resolved_max_batch_size(type(self)), + draft_model=getattr(self.llm, "draft_model", None), + facts=resolved_facts(type(self)), ) self._plan = plan self._mapping = getattr(model_config, "mapping", None) @@ -1192,6 +1325,10 @@ def __init__(self, model_config: ModelConfig, *args, **kwargs): # The built-in constructor returns early when it is re-entered; adopt once. if not isinstance(self.llm.model, MuseSparkOpenTepModelingV2): MuseSparkOpenTepModelingV2.adopt(self.llm.model, plan, capability) + if plan.levers is not None: + self.llm.model._arena_check = _arena_check( + getattr(self.llm, "spec_worker", None), plan.levers.deployment.arena + ) def load_weights(self, weights, weight_mapper=None): """The built-in load, then the plan's weight-dependent fast paths (the expert @@ -1293,4 +1430,9 @@ class derives from ``fused_accept``'s and ``draft_head``'s). """ from . import drafter as _drafter - _drafter.adopt_drafter(self, self._plan) + plan = self._plan + _drafter.adopt_drafter(self, plan) + if plan.levers is not None and plan.rank0 and not self.__dict__.get("_table_logged"): + # Once the drafter installs reported: the table every lever was installed with. + self._table_logged = True + logger.info(f"{_dispatch.TAG} lever table json (installed): {plan.report()}") diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py index 9181f8d25191..8380ab3fb463 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/weights.py @@ -63,16 +63,19 @@ def install_loaded_features( Then ``hot_replicas`` (``hot_replicas.install``: hooks only, at tp 4; the copies are gathered on the first forward of a deployment that serves 33..512-row decode steps), - then, on Rubin (sm107) at tp 4 (``prefill_cutedsl_experts.serves`` of - ``compute_capability``, the device's; None installs it nowhere), - ``prefill_cutedsl_experts`` (``prefill_cutedsl_experts.install``: the + then, with ``prefill_cutedsl_experts`` on Rubin (sm107) at tp 4 + (``prefill_cutedsl_experts.serves`` of ``compute_capability``, the device's; None + installs it nowhere), ``prefill_cutedsl_experts`` (``prefill_cutedsl_experts.install``: the 4096..8192-row expert calls of the NVFP4 layers on the Rubin CuTe DSL GEMMs), then ``early_expert_ids`` (``early_ids.install``, on the layers just installed), then ``pool_into_route`` (``DirectPool.install_sites``, whose routed kernel variants carry the placement and the early expert ids), each collective over the - TP group, and last ``next_qkv_prefetch``'s wrapper around every hooked layer's + TP group, then ``next_qkv_prefetch``'s wrapper around every hooked layer's expert forward (``prefetch.wrap_expert_forwards``, outermost, so it sees - whichever expert call the paths above installed). + whichever expert call the paths above installed), and last, with the lossy + ``context_fp8_dense`` lever, ``context_fp8_dense.install`` (FP8 copies of the dense + projections for the large context calls, at the lever's parameters, with + ``context_fp8_attn_gather``'s gather; collective). """ if plan.enabled("replicated_experts"): _install_ring_buddy(text_model, mapping, plan) @@ -83,13 +86,21 @@ def install_loaded_features( # After the owner-window install, whose expert and routing hooks it wraps. _hot_replicas.install(text_model, mapping) - if compute_capability is not None: + if plan.enabled("prefill_cutedsl_experts"): from .kernels.prefill_cutedsl_experts import install as _prefill_cutedsl - if _prefill_cutedsl.serves(compute_capability, mapping.tp_size): + serves = compute_capability is not None and _prefill_cutedsl.serves( + compute_capability, mapping.tp_size + ) + if serves: # After the expert paths above, whose expert forwards it wraps for the large # prefill calls only. _prefill_cutedsl.install(text_model, mapping) + plan.installed( + "prefill_cutedsl_experts", + serves, + f"written for sm107 at tp 4, not {compute_capability} at tp {mapping.tp_size}", + ) if plan.enabled("early_expert_ids"): _install_early_expert_ids(text_model, mapping) if plan.enabled("pool_into_route"): @@ -104,6 +115,26 @@ def install_loaded_features( from .kernels.next_qkv_prefetch import prefetch as _qkv_prefetch _qkv_prefetch.wrap_expert_forwards(text_model) + if plan.enabled("context_fp8_dense"): + from .kernels.context_fp8_dense import install as _context_fp8_dense + + # Lossy, on only by its knob; the dense projections' BF16 weights are untouched. + gather = plan.enabled("context_fp8_attn_gather") + names = _context_fp8_dense.install( + text_model, + mapping, + compute_capability, + min_rows=plan.param("context_fp8_dense", "MIN_ROWS"), + projections=plan.param("context_fp8_dense", "PROJECTIONS"), + gather=gather, + ) + plan.installed("context_fp8_dense", bool(names), "no projection got an MXFP8 copy") + if gather: + plan.installed( + "context_fp8_attn_gather", + "qkv" in names, + "the QKV projection keeps BF16 here, so its input is gathered in BF16", + ) def _install_ring_buddy(text_model, mapping, plan: "_dispatch.Plan") -> None: diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py index 28a3782e8e1c..f0f6fd765b9f 100644 --- a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_drafter_context.py @@ -6,7 +6,8 @@ geometry) over pages are bitwise their dense calls, the appended block included. * ``context_kv_fork``'s commit empties the dummy slot, as the worker's own commit does. * The block decode over a unified target/draft ``KVCacheManagerV2`` is bitwise the dense one. -* Over pages ``compose`` plans every sm103 TEP4 fast path but the window compaction. +* Over pages the only sm103 TEP4 fast path ``compose`` leaves out for a structural rule is + the window compaction. """ import importlib.util @@ -236,7 +237,8 @@ def test_the_block_decode_over_a_unified_kv_cache_is_the_dense_arena_one(wide): def test_a_paged_context_plans_every_fast_path_but_the_window_compaction(monkeypatch): """sm103 TEP4 with a TRTLLM drafter, TRTLLM_DFLASH_WINDOWED_CTX_ARENA set: a paged context is - never windowed, so the window compaction is the one served path the plan leaves out.""" + never windowed, so the window compaction is the one served path a structural rule leaves + out; every other path off is off by its lever rule at ``max_batch_size`` 64.""" monkeypatch.setattr(dispatch, "missing_preconditions", lambda *args: []) monkeypatch.setattr(dispatch, "agreed_union", lambda mapping, local: list(dict.fromkeys(local))) model_config = types.SimpleNamespace( @@ -246,8 +248,20 @@ def test_a_paged_context_plans_every_fast_path_but_the_window_compaction(monkeyp ) worker = types.SimpleNamespace(_dflash_attention_backend="TRTLLM", windowed_ctx_arena=True) plan = dispatch.compose( - None, model_config, ModelingV2Mode.REQUIRE, modeling.TARGET_PATHS, (10, 3), worker, 64 + None, + model_config, + ModelingV2Mode.REQUIRE, + modeling.TARGET_PATHS, + (10, 3), + worker, + 64, + snapshot={}, ) - served, _ = dispatch.capacity_gated(frozenset(modeling.TARGET_PATHS[103, 4].served), 64) - served, _ = dispatch.yielded(served) - assert served - plan.paths == {"drafter_window_compact"} + decisions = plan.levers.decisions + assert plan.levers.deployment.arena.value == "paged" + structural = {name for name, d in decisions.items() if d.source == "structural"} + assert structural == {"drafter_window_compact"} + off = {name for name, d in decisions.items() if not d.on} + # The half expert window's rule keeps it to max_batch_size 8; nothing builds on it. + assert off == {"drafter_window_compact", "half_window_experts"} + assert plan.paths == modeling.TARGET_PATHS[103, 4].served - off diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py new file mode 100644 index 000000000000..e22fad03d077 --- /dev/null +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py @@ -0,0 +1,1318 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The musespark_open TEP target's lever table, on the host only (no torch, no device). + +``levers.py`` and ``dispatch.py`` are loaded from their files, with the three leaf modules +``dispatch`` imports stubbed for the load (``_router_index``, the communicator and the +logger); the compositions are ``modeling.py``'s ``TARGET_PATHS`` statements run over that +``dispatch``. The tests pin: + +* the table of every region deployment under ``auto``, and of the replay knob sets (``R_LL``, ``R_H``, ``B_LL``); +* the ladder shapes, where an agg and a gen server resolve alike; +* ``baseline``: every key's plan is the composition before the lever table + (``C55_SERVED`` under ``c55_plan``'s rules) plus the two installs the plan adds to it; +* the knob grammar's errors, the overrides, the facts and the drafter arena; +* ``compose`` agreed over a stubbed TP group, and the ``Plan`` reports. +""" + +from __future__ import annotations + +import ast +import enum +import importlib +import importlib.util +import json +import re +import sys +import types +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +_ROOT = Path(__file__).resolve().parents[7] +_TARGET = ( + _ROOT / "tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open" + "/muse_spark_open_816b_a42b_nvfp4__sm_103__tep" +) +_PACKAGE = "_musespark_open_tep_levers_under_test" + + +class _Mode(str, enum.Enum): + OFF = "off" + AUTO = "auto" + REQUIRE = "require" + + +class _Log: + """The logger stub: every call as (level, message).""" + + def __init__(self) -> None: + self.records: list = [] + + def _record(self, level: str, message) -> None: + self.records.append((level, str(message))) + + def info(self, message, *args, **kwargs) -> None: + self._record("info", message) + + def warning(self, message, *args, **kwargs) -> None: + self._record("warning", message) + + def debug(self, message, *args, **kwargs) -> None: + self._record("debug", message) + + def error(self, message, *args, **kwargs) -> None: + self._record("error", message) + + def messages(self, level: str) -> list: + return [message for kind, message in self.records if kind == level] + + +class _Group: + """``Distributed.get``'s TP group: this rank's items, and each other rank's as + ``others`` makes them from this rank's (the same items by default).""" + + def __init__(self) -> None: + self.world = 1 + self.others = None + + def tp_allgather(self, items, small_payload: bool = False) -> list: + others = self.others or list + return [list(items)] + [others(list(items)) for _ in range(self.world - 1)] + + +class _Distributed: + group = _Group() + + @classmethod + def get(cls, mapping) -> _Group: + cls.group.world = int(mapping.tp_size) + return cls.group + + +_LOG = _Log() + + +def _module(name: str, **attrs) -> types.ModuleType: + module = types.ModuleType(name) + module.__dict__.update(attrs) + return module + + +def _load(): + stubs = { + "tensorrt_llm._torch._experimental.modeling_v2._router_index": _module( + "_router_index", ModelingV2Mode=_Mode + ), + "tensorrt_llm._torch.distributed.communicator": _module( + "communicator", Distributed=_Distributed + ), + "tensorrt_llm.logger": _module("logger", logger=_LOG), + } + saved = {name: sys.modules.get(name) for name in stubs} + sys.modules.update(stubs) + try: + package = types.ModuleType(_PACKAGE) + package.__path__ = [str(_TARGET)] + sys.modules[_PACKAGE] = package + levers = importlib.import_module(f"{_PACKAGE}.levers") + dispatch = importlib.import_module(f"{_PACKAGE}.dispatch") + finally: + for name, module in saved.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + return levers, dispatch + + +def _assigned(node) -> list: + if isinstance(node, ast.Assign): + return [target.id for target in node.targets if isinstance(target, ast.Name)] + if isinstance(node, (ast.AnnAssign, ast.AugAssign)) and isinstance(node.target, ast.Name): + return [node.target.id] + return [] + + +def _compositions(dispatch): + """``modeling.py``'s composition tables (its ``_TEP*``, ``_SM10*``, ``_PREFETCH*`` and + ``TARGET_PATHS`` statements), run over ``dispatch``.""" + path = _TARGET / "modeling.py" + wanted = re.compile(r"^(_(TEP|SM10|PREFETCH)\w*|TARGET_PATHS)$") + body = [ + node + for node in ast.parse(path.read_text()).body + if _assigned(node) and all(wanted.match(name) for name in _assigned(node)) + ] + namespace = {"_dispatch": dispatch} + exec(compile(ast.Module(body=body, type_ignores=[]), str(path), "exec"), namespace) + return namespace["TARGET_PATHS"] + + +LEVERS, DISPATCH = _load() +TARGET_PATHS = _compositions(DISPATCH) +RULES = LEVERS.RULES +BASELINE = {LEVERS.DEFAULT_KNOB: "baseline"} + + +def knobs(**values) -> dict: + """Lever knobs: ``name="on"``, or ``name__KEY="value"`` for a parameter.""" + out = {} + for spec, value in values.items(): + name, _, key = spec.partition("__") + out[LEVERS.knob_name(name, key or None)] = str(value) + return out + + +def table(key, mbs, role, arena, knob_set=None, facts=None): + sm, tp = key + return DISPATCH.table_for(TARGET_PATHS, sm, tp, mbs, role, arena, knob_set, facts) + + +def rule_levers_on(lever_table) -> set: + return {name for name in RULES if lever_table.on(name)} + + +def base_off(lever_table) -> set: + return { + name + for name, decision in lever_table.decisions.items() + if name not in RULES and not decision.on + } + + +# -------------------------------------------------------------------------- +# The registry +# -------------------------------------------------------------------------- +#: Where each rule-table lever is served, and its kind and origin. +SERVED_AT = { + "packed_prep_append": ({(107, 4)}, "lossless-bitwise", "HC"), + "windowed_commit_parallel": ({(107, 4), (107, 8)}, "lossless-bitwise", "HC"), + "context_fp8_dense": ({(107, 4)}, "lossy", "HC"), + "context_fp8_attn_gather": ({(107, 4)}, "lossy", "HC"), + "wide_direct_transports": ({(107, 4)}, "lossless-bitwise", "midc"), + "wide_router_per_row": ({(107, 4)}, "lossless-numeric", "midc"), + "wide_oproj_publish": ({(107, 4)}, "lossless-bitwise", "midc"), + "wide_finalize_transport": ({(107, 4)}, "lossless-bitwise", "midc"), + "lm_head_inbox": ({(103, 4), (103, 8), (107, 8)}, "composition", "upstream"), + "graph_start_qkv_prefetch": ({(103, 8), (107, 4), (107, 8)}, "composition", "upstream"), + "step_loop_on_execution_stream": ({(107, 4), (107, 8)}, "lossless-bitwise", "LL"), + "incremental_block_offsets": ({(107, 4), (107, 8)}, "lossless-bitwise", "LL"), + "drafter_kv_side_chain": ({(107, 8)}, "lossless-bitwise", "LL"), + "drafter_fc_early_fork": ({(107, 8)}, "lossless-bitwise", "LL"), + "drafter_preamble_from_commit": ({(107, 8)}, "lossless-bitwise", "LL"), + "head_sampler_coop_fold": ({(107, 4), (107, 8)}, "lossless-bitwise", "LL"), + "head_multimem_gather": ({(107, 8)}, "flagged", "LL"), + "half_window_experts": ({(103, 4), (107, 4)}, "composition", "upstream"), + "hot_replicas": ({(103, 4), (107, 4)}, "composition", "upstream"), + "prefill_cutedsl_experts": ({(107, 4)}, "composition", "upstream"), +} + + +def test_the_rule_table_is_the_registry_of_twenty_levers(): + assert set(RULES) == set(SERVED_AT) + assert set(RULES) <= DISPATCH.FAST_PATH_NAMES + for name, (keys, kind, origin) in SERVED_AT.items(): + served = {key for key, paths in TARGET_PATHS.items() if name in paths.served} + assert served == keys, name + assert (RULES[name].kind.value, RULES[name].origin) == (kind, origin), name + + +def test_no_lever_of_another_origin_is_served_on_sm103(): + others = {name for name, rule in RULES.items() if rule.origin != "upstream"} + for key in ((103, 4), (103, 8)): + assert not others & TARGET_PATHS[key].served + + +# -------------------------------------------------------------------------- +# The regions under auto +# -------------------------------------------------------------------------- +LL = {"step_loop_on_execution_stream", "incremental_block_offsets", "head_sampler_coop_fold"} +DENSE = {"drafter_kv_side_chain", "drafter_fc_early_fork", "drafter_preamble_from_commit"} +MIDC = {"wide_direct_transports", "wide_oproj_publish", "wide_finalize_transport"} +HC = {"packed_prep_append", "windowed_commit_parallel"} +PREFILL = {"prefill_cutedsl_experts"} + +REGIONS = [ + pytest.param((107, 8), "agg", (1,), "dense", LL | DENSE, id="ll-agg-tep8-mbs1"), + pytest.param( + (107, 4), + "agg", + (1,), + "windowed", + LL | {"half_window_experts", "wide_finalize_transport"} | PREFILL, + id="ll-agg-tep4-mbs1", + ), + pytest.param((107, 8), "gen", (1,), "windowed", LL, id="gen-tep8-b1"), + pytest.param( + (107, 4), + "gen", + (1,), + "windowed", + LL | {"half_window_experts", "wide_finalize_transport"} | PREFILL, + id="gen-tep4-b1", + ), + pytest.param( + (107, 4), + "gen", + (2, 3, 4), + "windowed", + MIDC | {"wide_router_per_row"} | LL | {"half_window_experts"} | PREFILL, + id="gen-tep4-b2-b4", + ), + pytest.param( + (107, 4), + "agg", + (2, 3, 4), + "windowed", + MIDC | {"wide_router_per_row"} | LL | {"half_window_experts"} | PREFILL, + id="agg-caps-2-4", + ), + pytest.param( + (107, 4), + "agg", + (6, 8), + "windowed", + HC + | MIDC + | {"incremental_block_offsets", "head_sampler_coop_fold", "half_window_experts"} + | PREFILL, + id="agg-caps-6-8", + ), + pytest.param( + (107, 4), + "agg", + (16, 32), + "windowed", + HC | MIDC | {"head_sampler_coop_fold", "hot_replicas"} | PREFILL, + id="agg-mid-high", + ), + pytest.param( + (107, 4), + "agg", + (48,), + "windowed", + HC | MIDC | {"head_sampler_coop_fold", "hot_replicas"} | PREFILL, + id="agg-hc", + ), + pytest.param( + (107, 4), + "gen", + (8,), + "windowed", + HC + | MIDC + | {"incremental_block_offsets", "head_sampler_coop_fold", "half_window_experts"} + | PREFILL, + id="gen-tep4-b8", + ), + pytest.param( + (107, 4), + "gen", + (16, 32, 64), + "windowed", + HC | MIDC | {"head_sampler_coop_fold", "hot_replicas"} | PREFILL, + id="gen-tep4-b16-b64", + ), + pytest.param((107, 8), "agg", (2, 4), "dense", LL | DENSE, id="agg-tep8-mbs2-4"), + pytest.param( + (107, 8), + "agg", + (8,), + "dense", + {"incremental_block_offsets", "head_sampler_coop_fold"} | DENSE, + id="agg-tep8-mbs8", + ), + pytest.param( + (107, 4), "ctx", (8, 64), "windowed", {"wide_finalize_transport"} | PREFILL, id="ctx-tep4" + ), + pytest.param( + (103, 8), + "agg", + (1, 8), + "dense", + {"lm_head_inbox", "graph_start_qkv_prefetch"}, + id="sm103-tep8", + ), + pytest.param( + (103, 4), + "agg", + (1, 4, 8), + "windowed", + {"lm_head_inbox", "half_window_experts"}, + id="sm103-tep4-low", + ), + pytest.param( + (103, 4), "agg", (16, 64), "windowed", {"lm_head_inbox", "hot_replicas"}, id="sm103-tep4" + ), +] + + +@pytest.mark.parametrize("key, role, sizes, arena, expected", REGIONS) +def test_region_tables_under_auto(key, role, sizes, arena, expected): + for mbs in sizes: + lever_table = table(key, mbs, role, arena) + assert rule_levers_on(lever_table) == expected, mbs + assert base_off(lever_table) == set(), mbs + assert lever_table.paths() == TARGET_PATHS[key].served - { + name for name, d in lever_table.decisions.items() if not d.on + } + + +def test_the_router_per_row_gates(): + on = table((107, 4), 4, "gen", "windowed").decisions["wide_router_per_row"] + assert on.on and on.source == "auto" and "gate" in on.why + off = table((107, 4), 8, "agg", "windowed").decisions["wide_router_per_row"] + assert not off.on and off.source == "gate" + assert not table((107, 4), 1, "agg", "windowed").on("wide_router_per_row") + + +def test_the_lossy_and_flagged_levers_are_never_on_under_auto(): + for key, paths in TARGET_PATHS.items(): + for mbs in (1, 4, 8, 64): + lever_table = table(key, mbs, "agg", "dense") + for name, rule in RULES.items(): + if rule.kind.value in ("lossy", "flagged") and name in paths.served: + assert not lever_table.on(name), (key, mbs, name) + + +@pytest.mark.parametrize("key", [(103, 4), (103, 8)]) +def test_sm103_resolves_under_auto_as_under_baseline(key): + for role in ("agg", "gen"): + for arena in ("dense", "windowed", "paged"): + for mbs in (None, 1, 4, 8, 9, 64): + auto = table(key, mbs, role, arena) + baseline = table(key, mbs, role, arena, BASELINE) + assert auto.paths() == baseline.paths(), (role, arena, mbs) + + +# -------------------------------------------------------------------------- +# The replay knob sets and the ladders +# -------------------------------------------------------------------------- +R_LL = knobs(head_multimem_gather="on") +R_H = knobs( + wide_direct_transports="off", + wide_oproj_publish="off", + wide_finalize_transport="off", + head_sampler_coop_fold="off", + graph_start_qkv_prefetch="on", + windowed_commit_parallel__MIN_REQUESTS=1, + context_fp8_dense="on", + context_fp8_attn_gather="on", +) +B_LL = knobs( + step_loop_on_execution_stream="off", + incremental_block_offsets="off", + head_sampler_coop_fold="off", + graph_start_qkv_prefetch="on", +) + + +def test_r_ll_selects_the_low_latency_set(): + lever_table = table((107, 8), 1, "agg", "dense", R_LL) + assert rule_levers_on(lever_table) == LL | DENSE | {"head_multimem_gather"} + assert lever_table.decisions["head_multimem_gather"].source == "knob" + assert not lever_table.on("lm_head_inbox") and not lever_table.on("graph_start_qkv_prefetch") + + +def test_r_h_selects_the_high_concurrency_set(): + lever_table = table((107, 4), 48, "agg", "windowed", R_H) + assert ( + rule_levers_on(lever_table) + == HC + | { + "context_fp8_dense", + "context_fp8_attn_gather", + "graph_start_qkv_prefetch", + "hot_replicas", + } + | PREFILL + ) + assert lever_table.param("windowed_commit_parallel", "MIN_REQUESTS") == 1 + assert lever_table.param("packed_prep_append", "MIN_ROWS") == 33 + assert lever_table.param("context_fp8_dense", "MIN_ROWS") == 1024 + assert lever_table.param("context_fp8_dense", "PROJECTIONS") == ( + "qkv", + "o_proj", + "pool", + "unpool", + ) + + +@pytest.mark.parametrize("mbs", [1, 2, 3, 4]) +def test_b_ll_selects_the_generation_set(mbs): + lever_table = table((107, 4), mbs, "gen", "windowed", B_LL) + wide = ( + {"wide_direct_transports", "wide_router_per_row", "wide_oproj_publish"} + if mbs > 1 + else set() + ) + assert ( + rule_levers_on(lever_table) + == wide + | { + "wide_finalize_transport", + "graph_start_qkv_prefetch", + "half_window_experts", + } + | PREFILL + ) + + +def test_baseline_tep8_gen_mbs1_is_the_previous_generation_selection(): + lever_table = table((107, 8), 1, "gen", "windowed", BASELINE) + assert rule_levers_on(lever_table) == {"lm_head_inbox", "graph_start_qkv_prefetch"} + + +R_B1 = dict(B_LL, **knobs(wide_finalize_transport="off")) +_LL_REPLAY = LL | DENSE | {"head_multimem_gather"} + + +@pytest.mark.parametrize( + "key, mbs, role, arena, knob_set, extra, less", + [ + # model_boundary has no knob: the replayed code built layer 0's input in the forward. + pytest.param( + (107, 8), 1, "agg", "dense", R_LL, _LL_REPLAY, {"lm_head_inbox"}, id="ll-tep8-mbs1" + ), + pytest.param( + (107, 4), + 48, + "agg", + "windowed", + R_H, + HC + | PREFILL + | {"context_fp8_dense", "context_fp8_attn_gather", "graph_start_qkv_prefetch"}, + set(), + id="hc-tep4-mbs48", + ), + *[ + pytest.param( + (107, 4), + mbs, + "gen", + "windowed", + B_LL, + MIDC | PREFILL | {"wide_router_per_row", "graph_start_qkv_prefetch"}, + set(), + id=f"gen-tep4-mbs{mbs}", + ) + for mbs in (2, 3, 4) + ], + pytest.param( + (107, 4), + 1, + "gen", + "windowed", + R_B1, + PREFILL | {"graph_start_qkv_prefetch"}, + set(), + id="gen-tep4-mbs1", + ), + pytest.param( + (107, 8), + 1, + "gen", + "windowed", + BASELINE, + {"graph_start_qkv_prefetch"}, + set(), + id="baseline-tep8-gen-mbs1", + ), + ], +) +def test_a_replay_set_is_the_whole_plan_it_replays(key, mbs, role, arena, knob_set, extra, less): + want = (c55_plan(key, mbs, arena) | extra) - less + assert set(table(key, mbs, role, arena, knob_set).paths()) == want + + +def test_auto_differs_from_r_h_only_by_the_small_step_levers(): + opt_ins = knobs(context_fp8_dense="on", context_fp8_attn_gather="on") + auto = table((107, 4), 48, "agg", "windowed", opt_ins).decisions + replay = table((107, 4), 48, "agg", "windowed", R_H).decisions + differ = { + name + for name, d in auto.items() + if (d.on, dict(d.params)) != (replay[name].on, dict(replay[name].params)) + } + assert differ == { + "graph_start_qkv_prefetch", + "head_sampler_coop_fold", + "wide_direct_transports", + "wide_finalize_transport", + "wide_oproj_publish", + "windowed_commit_parallel", + } + + +@pytest.mark.parametrize("arm", [None, B_LL, BASELINE], ids=["auto", "b-ll", "baseline"]) +@pytest.mark.parametrize("key, mbs", [((107, 4), 4), ((107, 4), 1), ((107, 8), 1)]) +def test_the_ladder_shapes_resolve_alike_on_agg_and_gen(key, mbs, arm): + def selection(role): + lever_table = table(key, mbs, role, "windowed", arm) + return {n: (d.on, dict(d.params)) for n, d in lever_table.decisions.items()} + + assert selection("agg") == selection("gen") + + +# -------------------------------------------------------------------------- +# baseline: the composition before the lever table +# -------------------------------------------------------------------------- +#: The served sets of the compositions at c55183f445 (its ``modeling.py`` tables over its +#: ``dispatch.FAST_PATH_NAMES``). +C55_SERVED = { + (103, 4): frozenset( + { + "batch_dense_gemms", + "bf16_routed_experts", + "boundary_early_release", + "capture_fold", + "context_kv_fork", + "context_prep", + "context_sequence_parallel", + "decode_kv_splits", + "deferred_input_prep", + "dense_gemms", + "draft_head", + "drafter_block_conv", + "drafter_block_decode", + "drafter_fc_shard", + "drafter_fixed_gemms", + "drafter_hops", + "drafter_pdl_chain", + "drafter_window_compact", + "finalize_transport", + "finalize_two_shot", + "full_attention_prep_append", + "full_attention_prep_fold", + "full_attention_split_kv", + "fused_accept", + "half_window_experts", + "hot_replicas", + "lm_head_inbox", + "model_boundary", + "moe_prologue", + "multi_request_post_attention_hop", + "multi_request_prep_append", + "multi_request_split_kv", + "path_walk_tail", + "persistent_experts", + "post_attention_hop", + "preattn_prep", + "projection_gathers", + "projection_publish", + "replicated_embedding", + "router_gemm", + "sampler_store", + "split_sampler", + "swa_block", + "swa_paged_decode", + "verify_fold", + "vocab_parallel_head", + "wide_draft_topk", + "wide_drafter_hops", + "wide_drafter_window_attention", + "wide_moe_allreduce", + "wide_projection_gathers", + "wide_split_sampler", + } + ), + (103, 8): frozenset( + { + "attention_output_into_hop", + "bf16_routed_experts", + "boundary_early_release", + "capture_fold", + "context_kv_fork", + "deferred_input_prep", + "dense_gemm_prefetch", + "dense_gemms", + "draft_head", + "drafter_block_conv", + "drafter_block_decode", + "drafter_fc_shard", + "drafter_fixed_gemms", + "drafter_hops", + "drafter_pdl_chain", + "early_expert_ids", + "finalize_transport", + "full_attention_prep_append", + "full_attention_prep_fold", + "full_attention_split_kv", + "fused_accept", + "lm_head_inbox", + "model_boundary", + "moe_prologue", + "next_qkv_prefetch", + "path_walk_tail", + "persistent_experts", + "pool_into_route", + "post_attention_hop", + "preattn_prep", + "projection_gathers", + "replicated_embedding", + "replicated_experts", + "sampler_store", + "split_sampler", + "swa_block", + "swa_paged_decode", + "three_window_placement", + "unpool_into_boundary", + "verify_fold", + } + ), + (107, 4): frozenset( + { + "attention_output_into_hop", + "batch_dense_gemms", + "bf16_routed_experts", + "boundary_early_release", + "capture_fold", + "context_kv_fork", + "context_prep", + "context_sequence_parallel", + "deferred_input_prep", + "dense_gemm_prefetch", + "dense_gemms", + "draft_head", + "drafter_block_conv", + "drafter_block_decode", + "drafter_fc_inbox", + "drafter_fc_shard", + "drafter_fixed_gemms", + "drafter_hops", + "drafter_pdl_chain", + "early_expert_ids", + "finalize_transport", + "finalize_two_shot", + "full_attention_prep_append", + "full_attention_prep_fold", + "full_attention_split_kv", + "fused_accept", + "half_window_experts", + "hot_replicas", + "mixed_step_gen_rows", + "model_boundary", + "moe_prologue", + "multi_request_post_attention_hop", + "multi_request_prep_append", + "multi_request_split_kv", + "next_qkv_prefetch", + "path_walk_tail", + "persistent_experts", + "pool_into_route", + "post_attention_hop", + "preattn_prep", + "projection_gathers", + "projection_publish", + "replicated_embedding", + "router_gemm", + "sampler_store", + "split_sampler", + "swa_block", + "swa_paged_decode", + "unpool_into_boundary", + "verify_fold", + "wide_dense_gemms", + "wide_draft_topk", + "wide_drafter_hops", + "wide_drafter_window_attention", + "wide_full_attention", + "wide_moe_allreduce", + "wide_persistent_experts", + "wide_post_attention_hop", + "wide_prologue", + "wide_split_sampler", + "wide_swa_block", + } + ), + (107, 8): frozenset( + { + "attention_output_into_hop", + "bf16_routed_experts", + "boundary_early_release", + "capture_fold", + "context_kv_fork", + "deferred_input_prep", + "dense_gemm_prefetch", + "dense_gemms", + "draft_head", + "drafter_block_conv", + "drafter_block_decode", + "drafter_fc_shard", + "drafter_fixed_gemms", + "drafter_hops", + "drafter_pdl_chain", + "early_expert_ids", + "finalize_transport", + "full_attention_prep_append", + "full_attention_prep_fold", + "full_attention_split_kv", + "fused_accept", + "lm_head_inbox", + "model_boundary", + "moe_prologue", + "next_qkv_prefetch", + "path_walk_tail", + "persistent_experts", + "pool_into_route", + "post_attention_hop", + "preattn_prep", + "projection_gathers", + "replicated_embedding", + "replicated_experts", + "sampler_store", + "split_sampler", + "swa_block", + "swa_paged_decode", + "three_window_placement", + "unpool_into_boundary", + "verify_fold", + } + ), +} + + +def c55_plan(key, max_batch_size, arena) -> set: + """``C55_SERVED[key]`` under that composition's rules: ``half_window_experts`` only up to + ``max_batch_size`` 8 (left out for an unknown one), ``hot_replicas`` left out of a plan + that carries the half window, and the drafter arena's: over a windowed arena on an + architecture other than sm107, ``context_kv_fork`` and ``verify_fold`` only with the + window compaction; any other arena leaves the compaction out. No served path builds on + one of these names there, so each rule removes the name alone.""" + paths = set(C55_SERVED[key]) + if max_batch_size is None or max_batch_size > 8: + paths.discard("half_window_experts") + if "half_window_experts" in paths: + paths.discard("hot_replicas") + if arena == "windowed": + if "drafter_window_compact" not in paths and key[0] != 107: + paths -= {"context_kv_fork", "verify_fold"} + else: + paths.discard("drafter_window_compact") + return paths + + +@pytest.mark.parametrize("key", sorted(C55_SERVED)) +@pytest.mark.parametrize("arena", ["dense", "windowed", "paged"]) +@pytest.mark.parametrize("role", ["agg", "gen", "ctx"]) +def test_baseline_is_the_previous_composition(key, arena, role): + for mbs in (None, 1, 2, 4, 8, 9, 16, 64): + want = c55_plan(key, mbs, arena) + # The two installs the plan adds to c55_plan: the graph-start fork of next_qkv_prefetch, + # and the large prefill expert calls the sm107 TEP4 load installed whenever composed. + if "next_qkv_prefetch" in want: + want.add("graph_start_qkv_prefetch") + if key == (107, 4): + want.add("prefill_cutedsl_experts") + assert table(key, mbs, role, arena, BASELINE).paths() == want, mbs + + +def test_baseline_leaves_every_other_origin_off_unless_its_knob_sets_it(): + lever_table = table( + (107, 4), 4, "agg", "windowed", {**BASELINE, **knobs(head_sampler_coop_fold="on")} + ) + others = {name for name, rule in RULES.items() if rule.origin != "upstream"} + assert rule_levers_on(lever_table) & others == {"head_sampler_coop_fold"} + decision = lever_table.decisions["wide_direct_transports"] + assert decision.source == "baseline" and "midc" in decision.why + + +# -------------------------------------------------------------------------- +# The knob grammar +# -------------------------------------------------------------------------- +@pytest.mark.parametrize( + "snapshot, message", + [ + (knobs(packed_prep_apend="on"), "did you mean MUSESPARK_OPEN_LEVER_PACKED_PREP_APPEND"), + (knobs(step_loop_on_execution_stream="yes"), "expected one of ['auto', 'on', 'off']"), + ({LEVERS.DEFAULT_KNOB: "fast"}, "expected one of ['auto', 'baseline']"), + ({LEVERS.ROLE_KNOB: "decode"}, "MUSESPARK_OPEN_LEVER__ROLE='decode'"), + (knobs(packed_prep_append__MIN_ROWS=7), "expected an integer in 8..1024, got 7"), + (knobs(packed_prep_append__MIN_ROWS="many"), "expected an integer in 8..1024"), + (knobs(packed_prep_append__MAX_ROWS=8), "declares the parameters ['MIN_ROWS']"), + (knobs(drafter_kv_side_chain__ROWS=8), "declares no parameter"), + (knobs(context_fp8_dense__MIN_ROWS=512), "expected an integer in 1024..32768"), + (knobs(context_fp8_dense__PROJECTIONS="qkv,lm_head"), "comma-separated subset"), + (knobs(head_sampler_coop_fold__MAX_ROWS=129), "expected an integer in 1..128, got 129"), + (knobs(head_sampler_coop_fold__MAX_ROWS=0), "expected an integer in 1..128, got 0"), + ( + knobs(windowed_commit_parallel__MIN_REQUESTS=65), + "expected an integer in 1..64, got 65", + ), + ( + knobs(step_loop_on_execution_stream__MAX_BATCH=65), + "expected an integer in 1..64, got 65", + ), + (knobs(incremental_block_offsets__MAX_BATCH=65), "expected an integer in 1..64, got 65"), + (knobs(half_window_experts__MAX_BATCH=65), "expected an integer in 1..64, got 65"), + ({"MUSESPARK_OPEN_LEVER_Packed_Prep_Append": "on"}, "no lever of this target"), + ], +) +def test_a_malformed_knob_raises(snapshot, message): + with pytest.raises(ValueError, match=re.escape(message)): + table((107, 4), 4, "agg", "windowed", snapshot) + + +def test_every_error_is_named_in_one_message(): + snapshot = {**knobs(packed_prep_apend="on", step_loop_on_execution_stream="yes")} + snapshot["MUSESPARK_OPEN_PS_SIDE_CHAIN"] = "1" + with pytest.raises(ValueError) as raised: + table((107, 4), 4, "agg", "windowed", snapshot) + text = str(raised.value) + assert "PACKED_PREP_APEND" in text and "'yes'" in text and "is retired" in text + + +#: The retired switches, pinned here so that a name dropped from ``levers.RETIRED`` fails. +SPEC_RETIRED = ( + "MUSESPARK_OPEN_PREP_PACKED_ROWS", + "MUSESPARK_OPEN_CTX_COMMIT_TUNED", + "MUSESPARK_OPEN_CONTEXT_FP8_DENSE", + "MUSESPARK_OPEN_CONTEXT_FP8_DENSE_MIN_ROWS", + "MUSESPARK_OPEN_CONTEXT_FP8_DENSE_PROJECTIONS", + "MUSESPARK_OPEN_CONTEXT_FP8_DENSE_GATHER", + "TRTLLM_MUSESPARK_TEP_WIDE_DIRECT", + "TRTLLM_MUSESPARK_TEP_WIDE_ROUTER_ROWS", + "TRTLLM_MUSESPARK_TEP_WIDE_OPROJ_PUBLISH", + "MUSESPARK_OPEN_STEP_LOOP_ONE_STREAM", + "MUSESPARK_OPEN_INCREMENTAL_BLOCK_OFFSETS", + "MUSESPARK_OPEN_PS_SIDE_CHAIN", + "MUSESPARK_OPEN_PS_EARLY_FC", + "MUSESPARK_OPEN_PS_DRAFTER_PREAMBLE", + "MUSESPARK_OPEN_PS_COOP_FOLD", + "MUSESPARK_OPEN_PS_SYMM_GATHER", + "MUSESPARK_OPEN_PS_NO_START_PREFETCH", + "MUSESPARK_OPEN_PS_NO_LM_HEAD_INBOX", + "MUSESPARK_OPEN_PS_LAYER0_INPUT", + "MUSESPARK_OPEN_PS_STEP_COMMIT_PDL", + "MUSESPARK_OPEN_HEAD_INBOX_POLL_SCOPE", + "MUSESPARK_OPEN_HEAD_INBOX_RESET_SCOPE", + "TRTLLM_MUSESPARK_SWA_BLOCKS_PER_PROGRAM", + "TRTLLM_MUSESPARK_TEP_WIDE_FULL_FOLD", + "TRTLLM_MUSESPARK_TEP_WIDE_QKV_PREFETCH", + "TRTLLM_MUSESPARK_TEP_WIDE_CAPTURE_FOLD", +) + + +def test_retired_is_the_spec_list(): + assert set(LEVERS.RETIRED) == set(SPEC_RETIRED) + + +@pytest.mark.parametrize("name", sorted(SPEC_RETIRED) + ["MUSESPARK_OPEN_PS_ANYTHING_ELSE"]) +def test_a_retired_switch_raises_naming_its_replacement(name): + replacement = LEVERS.retired_replacement(name) + assert replacement + with pytest.raises(ValueError, match=re.escape(f"{name} is retired (use {replacement})")): + table((107, 4), 4, "agg", "windowed", {name: "0"}) + + +@pytest.mark.parametrize( + "key, mbs, arena, snapshot, message", + [ + ( + (103, 8), + 1, + "dense", + knobs(step_loop_on_execution_stream="on"), + "does not serve step_loop_on_execution_stream", + ), + ((107, 8), 1, "windowed", knobs(drafter_kv_side_chain="on"), "dense full-length"), + ( + (107, 8), + 1, + "dense", + knobs(drafter_kv_side_chain="off", drafter_fc_early_fork="on"), + "builds on drafter_kv_side_chain, which is off", + ), + ((107, 4), 4, "windowed", knobs(hot_replicas="on"), "cannot both be on"), + ( + (107, 8), + 8, + "dense", + knobs(head_multimem_gather="on"), + "requires step_loop_on_execution_stream, which is off", + ), + ( + (107, 8), + 1, + "dense", + knobs(head_multimem_gather="on", lm_head_inbox="on"), + "cannot both be on", + ), + ( + (107, 4), + 1, + "windowed", + knobs( + context_fp8_dense="on", + context_fp8_dense__PROJECTIONS="o_proj", + context_fp8_attn_gather="on", + ), + "PROJECTIONS leaves out", + ), + ], +) +def test_on_raises_where_the_lever_cannot_be_on(key, mbs, arena, snapshot, message): + with pytest.raises(ValueError, match=re.escape(message)): + table(key, mbs, "agg", arena, snapshot) + + +# -------------------------------------------------------------------------- +# Overrides +# -------------------------------------------------------------------------- +def test_off_turns_a_lever_and_what_builds_on_it_off(): + lever_table = table((107, 8), 1, "agg", "dense", knobs(drafter_kv_side_chain="off")) + decisions = lever_table.decisions + assert (decisions["drafter_kv_side_chain"].on, decisions["drafter_kv_side_chain"].source) == ( + False, + "knob", + ) + assert (decisions["drafter_fc_early_fork"].on, decisions["drafter_fc_early_fork"].source) == ( + False, + "dependency", + ) + assert decisions["drafter_preamble_from_commit"].on + + +@pytest.mark.parametrize("key", sorted(TARGET_PATHS)) +@pytest.mark.parametrize("arena", ["dense", "windowed", "paged"]) +def test_every_single_off_gives_a_closed_plan(key, arena): + for mbs in (1, 4, 8, 48): + for name in sorted(TARGET_PATHS[key].served): + paths = table(key, mbs, "agg", arena, knobs(**{name: "off"})).paths() + DISPATCH.Plan(paths) # raises when not closed under depends_on + for lever in paths & set(RULES): + assert set(RULES[lever].requires) <= paths, (key, mbs, name, lever) + + +def test_a_base_path_knob_turns_its_dependents_off(): + lever_table = table((107, 4), 8, "agg", "windowed", knobs(multi_request_prep_append="off")) + assert not lever_table.on("multi_request_prep_append") + assert lever_table.decisions["packed_prep_append"].source == "dependency" + + +def test_a_parameter_moves_the_band(): + lever_table = table((107, 4), 1, "agg", "windowed", knobs(packed_prep_append__MIN_ROWS=8)) + assert lever_table.on("packed_prep_append") + assert lever_table.param("packed_prep_append", "MIN_ROWS") == 8 + lever_table = table( + (107, 4), 1, "agg", "windowed", knobs(windowed_commit_parallel__MIN_REQUESTS=1) + ) + assert lever_table.on("windowed_commit_parallel") + lever_table = table( + (107, 8), 8, "agg", "dense", knobs(step_loop_on_execution_stream__MAX_BATCH=8) + ) + assert lever_table.on("step_loop_on_execution_stream") + + +def test_on_outside_a_numeric_gate_warns(): + lever_table = table((107, 4), 8, "agg", "windowed", knobs(wide_router_per_row="on")) + assert lever_table.on("wide_router_per_row") + assert any("outside its accuracy gates" in warning for warning in lever_table.warnings) + + +def test_a_knob_of_an_unserved_lever_set_off_or_auto_only_warns(): + lever_table = table((103, 8), 1, "agg", "dense", knobs(head_sampler_coop_fold="off")) + assert any("has no effect" in warning for warning in lever_table.warnings) + assert lever_table.paths() == table((103, 8), 1, "agg", "dense").paths() + + +def test_an_empty_knob_is_unset(): + snapshot = LEVERS.snapshot_of({LEVERS.knob_name("hot_replicas"): " ", "PATH": "/bin"}) + assert snapshot == {} + + +# -------------------------------------------------------------------------- +# Facts and the drafter arena +# -------------------------------------------------------------------------- +def test_the_facts_of_a_snapshot(): + snapshot = LEVERS.snapshot_of( + {LEVERS.KV_DEFER_CLOSE_ENV: "1", LEVERS.UCX_RNDV_THRESH_ENV: "8192", "HOME": "/root"} + ) + assert set(snapshot) == {LEVERS.KV_DEFER_CLOSE_ENV, LEVERS.UCX_RNDV_THRESH_ENV} + assert LEVERS.facts_of(snapshot) == LEVERS.Facts(kv_defer_close=True, ucx_rndv_thresh_set=True) + assert LEVERS.facts_of({LEVERS.KV_DEFER_CLOSE_ENV: "0"}) == LEVERS.Facts() + + +def test_the_deferred_close_leaves_the_one_stream_loop_out(): + facts = LEVERS.Facts(kv_defer_close=True) + lever_table = table((107, 8), 1, "agg", "dense", facts=facts) + decision = lever_table.decisions["step_loop_on_execution_stream"] + assert (decision.on, decision.source) == (False, "structural") + with pytest.raises(ValueError, match="requires step_loop_on_execution_stream"): + table((107, 8), 1, "agg", "dense", knobs(head_multimem_gather="on"), facts) + + +def test_rndv_raises_on_a_disaggregated_server_only(): + facts = LEVERS.Facts(ucx_rndv_thresh_set=True) + assert table((107, 4), 4, "agg", "windowed", facts=facts).deployment.facts == facts + for role in ("gen", "ctx"): + with pytest.raises(ValueError, match=LEVERS.UCX_RNDV_THRESH_ENV): + table((107, 4), 4, role, "windowed", facts=facts) + + +def test_rndv_on_an_unhinted_disaggregated_server_raises(): + facts = LEVERS.Facts(ucx_rndv_thresh_set=True) + with pytest.raises(ValueError, match=LEVERS.UCX_RNDV_THRESH_ENV): + table((107, 4), 4, "unknown", "windowed", facts=facts) + lever_table = DISPATCH.table_for( + TARGET_PATHS, 107, 4, 4, "unknown", "windowed", facts=facts, role_source="no engine facts" + ) + assert lever_table.deployment.role is LEVERS.Role.UNKNOWN + + +def test_context_fp8_dense_on_refuses_a_reachable_decode_step(): + with pytest.raises(ValueError, match="decode steps of 1024"): + table((107, 4), 128, "agg", "windowed", knobs(context_fp8_dense="on")) + with pytest.raises(ValueError, match="max_batch_size unknown"): + table((107, 4), None, "agg", "windowed", knobs(context_fp8_dense="on")) + moved = knobs(context_fp8_dense="on", context_fp8_dense__MIN_ROWS=2048) + assert table((107, 4), 128, "agg", "windowed", moved).on("context_fp8_dense") + assert table((107, 4), 48, "agg", "windowed", knobs(context_fp8_dense="on")).on( + "context_fp8_dense" + ) + + +def test_a_paged_arena_leaves_the_arena_writing_levers_out(): + lever_table = table((107, 8), 1, "agg", "paged") + for name in DENSE | {"incremental_block_offsets", "windowed_commit_parallel"}: + assert lever_table.decisions[name].source == "structural", name + assert lever_table.on("step_loop_on_execution_stream") + assert lever_table.on("context_kv_fork") + + +def test_the_window_compaction_needs_a_windowed_arena_and_the_full_context_paths_need_it(): + assert table((103, 4), 1, "agg", "windowed").on("drafter_window_compact") + dense = table((103, 4), 1, "agg", "dense") + assert dense.decisions["drafter_window_compact"].source == "structural" + sm103_tep8 = table((103, 8), 1, "agg", "windowed") + assert sm103_tep8.decisions["context_kv_fork"].source == "structural" + assert sm103_tep8.decisions["verify_fold"].source == "structural" + assert table((107, 8), 1, "agg", "windowed").on("context_kv_fork") + # With the compaction served, a knob that turns it off takes the arena-writing paths along. + off = table((103, 4), 1, "agg", "windowed", knobs(drafter_window_compact="off")) + assert off.decisions["context_kv_fork"].source == "dependency" + + +def test_the_arena_of_the_worker_and_its_drafter(): + paged = ("TRTLLM", "FA4") + worker = types.SimpleNamespace(_dflash_attention_backend="AUTO", windowed_ctx_arena=True) + drafter = types.SimpleNamespace(dflash_attention_backend="TRTLLM") + assert LEVERS.arena_of(worker, drafter, paged) is LEVERS.Arena.PAGED + drafter = types.SimpleNamespace(dflash_attention_backend="FLASHINFER") + assert LEVERS.arena_of(worker, drafter, paged) is LEVERS.Arena.WINDOWED + drafter = types.SimpleNamespace(_paged_ctx_cache=True) + assert LEVERS.arena_of(worker, drafter, paged) is LEVERS.Arena.PAGED + dense = types.SimpleNamespace(_dflash_attention_backend="FLASHINFER", windowed_ctx_arena=False) + assert LEVERS.arena_of(dense, None, paged) is LEVERS.Arena.DENSE + + +def test_the_arena_the_worker_built(): + assert LEVERS.worker_arena(types.SimpleNamespace(_ctx_buf_inited=False)) is None + built = dict(_ctx_buf_inited=True, _ctx_paged=False, _ctx_base=None) + assert LEVERS.worker_arena(types.SimpleNamespace(**built)) is LEVERS.Arena.DENSE + built["_ctx_base"] = object() + assert LEVERS.worker_arena(types.SimpleNamespace(**built)) is LEVERS.Arena.WINDOWED + built["_ctx_paged"] = True + assert LEVERS.worker_arena(types.SimpleNamespace(**built)) is LEVERS.Arena.PAGED + + +def test_the_role(): + hint = {LEVERS.ROLE_KNOB: "gen"} + assert LEVERS.role_of(hint, True) == (LEVERS.Role.GEN, "hint") + assert LEVERS.role_of({}, False) == (LEVERS.Role.AGG, "no cache transceiver") + assert LEVERS.role_of({}, True)[0] is LEVERS.Role.UNKNOWN + assert LEVERS.role_of({}, None)[0] is LEVERS.Role.UNKNOWN + with pytest.raises(ValueError, match="with a cache transceiver"): + LEVERS.role_of({LEVERS.ROLE_KNOB: "agg"}, True) + with pytest.raises(ValueError, match="without one|without a cache transceiver"): + LEVERS.role_of(hint, False) + unknown = table((107, 4), 4, "unknown", "windowed") + assert any("role unknown" in warning for warning in unknown.warnings) + # Counted as a decoding server; no accuracy gate names an unknown role. + gen = rule_levers_on(table((107, 4), 4, "gen", "windowed")) + assert rule_levers_on(unknown) == gen - {"wide_router_per_row"} + assert unknown.decisions["wide_router_per_row"].source == "gate" + + +def test_the_table_lines_and_json(): + lever_table = table((107, 4), 4, "gen", "windowed", knobs(hot_replicas="off")) + lines = lever_table.describe() + assert lines[0].startswith("levers: sm107 tep4, gen (hint)") + resolved = json.loads(lines[-1].split("lever table json (resolved): ", 1)[1]) + assert resolved == json.loads(lever_table.as_json()) + assert resolved["levers"]["hot_replicas"] == {"resolved": "off", "source": "knob"} + assert resolved["knobs"] == {LEVERS.knob_name("hot_replicas"): "off"} + assert resolved["deployment"]["max_batch_size"] == 4 + installed = json.loads(lever_table.as_json({"wide_oproj_publish": False})) + assert installed["levers"]["wide_oproj_publish"]["installed"] is False + assert installed["levers"]["wide_direct_transports"]["installed"] is True + + +# -------------------------------------------------------------------------- +# compose over a stubbed TP group, and the plan +# -------------------------------------------------------------------------- +@pytest.fixture +def compose(monkeypatch): + group = _Distributed.group + monkeypatch.setattr(group, "others", None) + _LOG.records.clear() + monkeypatch.setattr( + DISPATCH, + "composition", + lambda model_config, compositions, capability: compositions.get( + (DISPATCH.sm_of(capability), model_config.mapping.tp_size) + ), + ) + monkeypatch.setattr(DISPATCH, "missing_preconditions", lambda *args: []) + monkeypatch.setattr(DISPATCH, "unserved_reason", lambda *args: "not served") + dflash = types.ModuleType("tensorrt_llm._torch.speculative.dflash") + dflash._PAGED_ATTENTION_BACKENDS = ("TRTLLM", "FA4") + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.speculative.dflash", dflash) + + def run(key, mbs, snapshot=None, disaggregated=False, arena="dense", tp_rank=0): + sm, tp = key + mapping = types.SimpleNamespace( + tp_size=tp, tp_rank=tp_rank, gpus_per_node=4, tp_group=list(range(tp)) + ) + worker = types.SimpleNamespace( + _dflash_attention_backend="TRTLLM" if arena == "paged" else "FLASHINFER", + windowed_ctx_arena=arena == "windowed", + ) + facts = types.SimpleNamespace( + max_batch_size=mbs, max_num_tokens=8192, disaggregated=disaggregated + ) + return DISPATCH.compose( + None, + types.SimpleNamespace(mapping=mapping), + _Mode.AUTO, + TARGET_PATHS, + (sm // 10, sm % 10), + worker, + mbs, + facts=facts, + snapshot=snapshot or {}, + ) + + return run + + +def test_compose_plans_the_table_and_logs_it_on_rank_0(compose): + plan = compose((107, 4), 4, arena="windowed") + assert plan.paths == table((107, 4), 4, "agg", "windowed").paths() + assert plan.levers.deployment.role_source == "no cache transceiver" + assert any("lever table json (resolved)" in m for m in _LOG.messages("info")) + _LOG.records.clear() + compose((107, 4), 4, arena="windowed", tp_rank=1) + info = _LOG.messages("info") + assert any("drafter context arena" in m for m in info) + assert not any("lever table json" in m for m in info) + + +def test_compose_raises_when_the_ranks_disagree_on_the_knobs(compose): + _Distributed.group.others = lambda items: [item.replace("=on", "=off") for item in items] + with pytest.raises(ValueError, match="disagree"): + compose((107, 8), 1, knobs(head_multimem_gather="on")) + + +def test_compose_raises_another_ranks_lever_error(compose): + _Distributed.group.others = lambda items: items + ["lever table: a rank's own error"] + with pytest.raises(ValueError, match="a rank's own error"): + compose((107, 8), 1) + + +def test_compose_raises_another_ranks_missing_precondition(compose): + _Distributed.group.others = lambda items: items + ["a precondition"] + with pytest.raises(ValueError, match="the fast paths need a precondition"): + compose((107, 8), 1) + + +def test_compose_raises_this_ranks_lever_error_after_the_agreement(compose): + with pytest.raises(ValueError, match="cannot both be on"): + compose((107, 4), 4, knobs(hot_replicas="on"), arena="windowed") + + +def test_compose_without_a_composition_still_checks_the_knobs(compose, monkeypatch): + monkeypatch.setattr(DISPATCH, "composition", lambda *args: None) + monkeypatch.setattr(DISPATCH, "require_served", lambda *args: None) + assert compose((107, 4), 4, knobs(hot_replicas="on")) is DISPATCH.NO_PATHS + assert any("no effect without a composition" in m for m in _LOG.messages("info")) + with pytest.raises(ValueError, match="is retired"): + compose((107, 4), 4, {"MUSESPARK_OPEN_PS_SIDE_CHAIN": "1"}) + + +def test_the_role_hint_on_a_disaggregated_server(compose): + plan = compose((107, 4), 4, {LEVERS.ROLE_KNOB: "gen"}, disaggregated=True, arena="windowed") + assert plan.levers.deployment.role is LEVERS.Role.GEN + with pytest.raises(ValueError, match="with a cache transceiver"): + compose((107, 4), 4, {LEVERS.ROLE_KNOB: "agg"}, disaggregated=True) + + +def test_compose_names_a_bad_role_hint_with_every_other_grammar_error(compose): + snapshot = {LEVERS.ROLE_KNOB: "decode", **knobs(packed_prep_apend="on")} + with pytest.raises(ValueError) as raised: + compose((107, 4), 4, snapshot, disaggregated=True, arena="windowed") + text = str(raised.value) + assert "'decode'" in text and "PACKED_PREP_APEND" in text + + +def _router_index(monkeypatch): + monkeypatch.setitem(sys.modules, "tensorrt_llm.logger", _module("logger", logger=_LOG)) + path = _ROOT / "tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py" + spec = importlib.util.spec_from_file_location("_router_index_under_test", path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize( + "transceiver, disaggregated", + [ + (None, False), + (types.SimpleNamespace(backend=None), False), + (types.SimpleNamespace(backend="NIXL"), True), + ], +) +def test_a_server_is_disaggregated_only_with_a_transceiver_backend( + monkeypatch, transceiver, disaggregated +): + args = types.SimpleNamespace( + max_batch_size=4, + max_num_tokens=8192, + cuda_graph_config=None, + cache_transceiver_config=transceiver, + ) + facts = _router_index(monkeypatch)._facts_of(args) + assert facts.disaggregated is disaggregated + + +def test_the_plan_reports(compose): + plan = compose((107, 4), 4, knobs(wide_oproj_publish="on"), arena="windowed") + assert plan.param("packed_prep_append", "MIN_ROWS") == 33 + plan.installed("wide_direct_transports", False, "no wide workspace") + assert any("wide_direct_transports is not installed" in m for m in _LOG.messages("warning")) + with pytest.raises(ValueError, match="WIDE_OPROJ_PUBLISH=on"): + plan.installed("wide_oproj_publish", False, "refused") + report = json.loads(plan.report()) + assert report["levers"]["wide_direct_transports"]["installed"] is False + assert report["levers"]["half_window_experts"]["installed"] is True + for _ in range(3): + plan.declined("incremental_block_offsets", "a reason", "8 rows") + assert plan.declines == {("incremental_block_offsets", "a reason"): 3} + assert sum("declined a step" in m for m in _LOG.messages("warning")) == 1 + _LOG.records.clear() + for _ in range(DISPATCH.DECLINE_REPORT_STEPS): + plan.step() + assert any("lever declines after" in m for m in _LOG.messages("info")) + + +def test_a_plan_built_from_names_takes_the_declared_parameters(): + plan = DISPATCH.plan_of("split_sampler") + assert plan.param("head_sampler_coop_fold", "MAX_ROWS") == 16 + assert plan.report() == "{}" + with pytest.raises(KeyError): + plan.param("split_sampler", "MAX_ROWS") From c1d97c495b29353170c72013fc5a49f10d9e9648 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 14:41:41 +0000 Subject: [PATCH 11/15] [None][fix] modeling_v2 musespark_open tep: router per row only on generation servers wide_router_per_row is a lossless-numeric lever: each verify row routes with its own one-request bits, so near-tie experts can resolve differently from the built-in wide router. Its accuracy gate now holds only on generation servers at TEP4 with max_batch_size 2..4; aggregated servers of the same batch range resolve it off under auto (source "gate"). The lever tests now expect it off at agg mbs 2..4 and check that an aggregated server built from a generation config differs from that generation server only in this lever. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../levers.py | 5 +---- .../test_modeling_v2_levers.py | 11 +++++++++-- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py index 0a9acf7af56c..b980ba94ebdb 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py @@ -514,10 +514,7 @@ def holds(d: Deployment) -> bool: "midc", Kind.NUMERIC, _wide_steps, - gates=( - _router_gate(Role.GEN, "generation servers at tep4 with max_batch_size 2..4"), - _router_gate(Role.AGG, "aggregated servers at tep4 with max_batch_size 2..4"), - ), + gates=(_router_gate(Role.GEN, "generation servers at tep4 with max_batch_size 2..4"),), ), "wide_oproj_publish": Rule("midc", Kind.BITWISE, _wide_steps), "wide_finalize_transport": Rule( diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py index e22fad03d077..8a27d61c425d 100644 --- a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py @@ -267,7 +267,7 @@ def test_no_lever_of_another_origin_is_served_on_sm103(): "agg", (2, 3, 4), "windowed", - MIDC | {"wide_router_per_row"} | LL | {"half_window_experts"} | PREFILL, + MIDC | LL | {"half_window_experts"} | PREFILL, id="agg-caps-2-4", ), pytest.param( @@ -366,6 +366,9 @@ def test_the_router_per_row_gates(): assert on.on and on.source == "auto" and "gate" in on.why off = table((107, 4), 8, "agg", "windowed").decisions["wide_router_per_row"] assert not off.on and off.source == "gate" + for mbs in (2, 3, 4): + off = table((107, 4), mbs, "agg", "windowed").decisions["wide_router_per_row"] + assert not off.on and off.source == "gate", mbs assert not table((107, 4), 1, "agg", "windowed").on("wide_router_per_row") @@ -555,7 +558,11 @@ def selection(role): lever_table = table(key, mbs, role, "windowed", arm) return {n: (d.on, dict(d.params)) for n, d in lever_table.decisions.items()} - assert selection("agg") == selection("gen") + agg, gen = selection("agg"), selection("gen") + # The router-per-row accuracy gate holds on generation servers only. + gated = {"wide_router_per_row"} if gen.get("wide_router_per_row", (False,))[0] else set() + assert {n for n in agg if agg[n] != gen[n]} == gated + assert all(not agg[n][0] for n in gated) # -------------------------------------------------------------------------- From 110d5c2771802eecf8b746d737f5d52a94866102 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 31 Oct 2026 15:10:51 +0000 Subject: [PATCH 12/15] [None][feat] modeling_v2 musespark_open tep: high-concurrency lever selection Aggregated sm107 TEP4 servers with max_batch_size of at least HC_MIN_BATCH (48) resolve under auto to the high-concurrency selection: the wide decode transports, the wide o_proj publish, the wide finalize transport and the cooperative head-sampler fold are off; the graph-start QKV prefetch is on; and the parallel windowed context commit starts from one request. A lever Param can now declare a per-deployment default (default_for); a knob still overrides it. Generation servers, TEP8 and smaller batches keep their selection. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- .../levers.py | 67 +++++++++++++++---- .../test_modeling_v2_levers.py | 43 +++++++----- 2 files changed, 80 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py index b980ba94ebdb..d1655248e84c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/levers.py @@ -61,6 +61,10 @@ ROWS_PER_REQUEST = DRAFT_LEN + 1 #: The fewest decode rows of a step of two requests, where the wide TP4 paths start. WIDE_MIN_ROWS = 2 * ROWS_PER_REQUEST +#: Aggregated sm107 TEP4 servers of at least this max_batch_size run the high-concurrency +#: selection: no wide decode paths or cooperative head fold, the graph-start QKV prefetch, +#: and the parallel windowed context commit from one request. +HC_MIN_BATCH = 48 #: The prefix of every lever knob. PREFIX = "MUSESPARK_OPEN_LEVER_" @@ -240,7 +244,8 @@ class Param: An integer in ``low..high``, or, with ``choices``, a non-empty subset of ``choices`` (comma-separated, kept in ``choices`` order). A lossy lever's range is a subset of - the values its accuracy gate covered. + the values its accuracy gate covered. ``default_for``, when set, gives the default + per deployment; the knob still wins over it. """ key: str @@ -248,6 +253,10 @@ class Param: low: int = 1 high: int = 1 << 30 choices: Tuple[str, ...] = () + default_for: Optional[Callable[["Deployment"], object]] = None + + def default_of(self, deployment: "Deployment") -> object: + return self.default if self.default_for is None else self.default_for(deployment) def parse(self, raw: str) -> object: """The value ``raw`` gives this parameter. @@ -367,6 +376,29 @@ def param(self, name: str, key: str) -> object: # The deployment rules # -------------------------------------------------------------------------- _NO_DECODE = "the server runs no decode step" +_HC = f"aggregated sm107 tep4 servers of max_batch_size {HC_MIN_BATCH} or more" + + +def _hc_selection(d: Deployment) -> bool: + """Whether ``d`` runs the high-concurrency selection (``HC_MIN_BATCH``).""" + return ( + d.sm == 107 + and d.tp == 4 + and d.role is Role.AGG + and d.max_batch_size is not None + and d.max_batch_size >= HC_MIN_BATCH + ) + + +def _off_at_hc(rule: "RuleFn") -> "RuleFn": + """``rule``, except off on the high-concurrency selection.""" + + def wrapped(d: Deployment, t: "Resolving", p: Mapping[str, object]) -> Tuple[bool, str]: + if _hc_selection(d): + return False, f"{_HC} run without it" + return rule(d, t, p) + + return wrapped def _reaching(rows: int, what: str) -> RuleFn: @@ -442,6 +474,17 @@ def rule(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, st return rule +def _graph_start_prefetch(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: + if d.sm != 107: + return True, "graph starts of split-K row counts" + if _hc_selection(d): + return True, f"graph starts of split-K row counts on {_HC}" + return False, ( + "on sm107 a captured decode graph starts without a fork, and layer 0's QKV weight is " + "not prefetched" + ) + + def _coop_fold(d: Deployment, t: Resolving, p: Mapping[str, object]) -> Tuple[bool, str]: if not d.decodes(): return False, _NO_DECODE @@ -491,7 +534,9 @@ def holds(d: Deployment) -> bool: "HC", Kind.BITWISE, _windowed_commit_parallel, - params=(Param("MIN_REQUESTS", 5, 1, 64),), + params=( + Param("MIN_REQUESTS", 5, 1, 64, default_for=lambda d: 1 if _hc_selection(d) else 5), + ), ), "context_fp8_dense": Rule( "HC", @@ -509,16 +554,18 @@ def holds(d: Deployment) -> bool: "context_fp8_attn_gather": Rule( "HC", Kind.LOSSY, _lossy, after=("context_fp8_dense",), excluded=_attn_gather_excluded ), - "wide_direct_transports": Rule("midc", Kind.BITWISE, _wide_steps), + "wide_direct_transports": Rule("midc", Kind.BITWISE, _off_at_hc(_wide_steps)), "wide_router_per_row": Rule( "midc", Kind.NUMERIC, _wide_steps, gates=(_router_gate(Role.GEN, "generation servers at tep4 with max_batch_size 2..4"),), ), - "wide_oproj_publish": Rule("midc", Kind.BITWISE, _wide_steps), + "wide_oproj_publish": Rule("midc", Kind.BITWISE, _off_at_hc(_wide_steps)), "wide_finalize_transport": Rule( - "midc", Kind.BITWISE, _always("17..32-token finalize calls on a 32-token transport") + "midc", + Kind.BITWISE, + _off_at_hc(_always("17..32-token finalize calls on a 32-token transport")), ), "lm_head_inbox": Rule( "upstream", @@ -531,11 +578,7 @@ def holds(d: Deployment) -> bool: "graph_start_qkv_prefetch": Rule( "upstream", Kind.COMPOSITION, - _not_on_sm107( - "on sm107 a captured decode graph starts without a fork, and layer 0's QKV " - "weight is not prefetched", - "graph starts of split-K row counts", - ), + _graph_start_prefetch, baseline=_always("graph starts of split-K row counts"), ), "step_loop_on_execution_stream": Rule( @@ -560,7 +603,7 @@ def holds(d: Deployment) -> bool: "head_sampler_coop_fold": Rule( "LL", Kind.BITWISE, - _coop_fold, + _off_at_hc(_coop_fold), params=(Param("MAX_ROWS", 16, 1, 128),), ), "head_multimem_gather": Rule( @@ -1015,7 +1058,7 @@ def resolve( rule = RULES.get(name) mode = knobs.modes.get(name) params = { - param.key: knobs.params.get((name, param.key), param.default) + param.key: knobs.params.get((name, param.key), param.default_of(d)) for param in (rule.params if rule is not None else ()) } decision = _decide( diff --git a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py index 8a27d61c425d..0d58ab2a3753 100644 --- a/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py +++ b/tests/unittest/_torch/modeling_v2/models/musespark_open/muse_spark_open_816b_a42b_nvfp4__sm_103__tep/test_modeling_v2_levers.py @@ -154,6 +154,7 @@ def _compositions(dispatch): LEVERS, DISPATCH = _load() TARGET_PATHS = _compositions(DISPATCH) RULES = LEVERS.RULES +HC_MIN_BATCH = LEVERS.HC_MIN_BATCH BASELINE = {LEVERS.DEFAULT_KNOB: "baseline"} @@ -292,9 +293,9 @@ def test_no_lever_of_another_origin_is_served_on_sm103(): pytest.param( (107, 4), "agg", - (48,), + (48, 64), "windowed", - HC | MIDC | {"head_sampler_coop_fold", "hot_replicas"} | PREFILL, + HC | {"graph_start_qkv_prefetch", "hot_replicas"} | PREFILL, id="agg-hc", ), pytest.param( @@ -532,23 +533,29 @@ def test_a_replay_set_is_the_whole_plan_it_replays(key, mbs, role, arena, knob_s assert set(table(key, mbs, role, arena, knob_set).paths()) == want -def test_auto_differs_from_r_h_only_by_the_small_step_levers(): +def test_auto_at_the_high_concurrency_selection_is_r_h(): opt_ins = knobs(context_fp8_dense="on", context_fp8_attn_gather="on") - auto = table((107, 4), 48, "agg", "windowed", opt_ins).decisions - replay = table((107, 4), 48, "agg", "windowed", R_H).decisions - differ = { - name - for name, d in auto.items() - if (d.on, dict(d.params)) != (replay[name].on, dict(replay[name].params)) - } - assert differ == { - "graph_start_qkv_prefetch", - "head_sampler_coop_fold", - "wide_direct_transports", - "wide_finalize_transport", - "wide_oproj_publish", - "windowed_commit_parallel", - } + for mbs in (HC_MIN_BATCH, 64): + auto = table((107, 4), mbs, "agg", "windowed", opt_ins).decisions + replay = table((107, 4), mbs, "agg", "windowed", R_H).decisions + assert { + name + for name, d in auto.items() + if (d.on, dict(d.params)) != (replay[name].on, dict(replay[name].params)) + } == set(), mbs + + +def test_the_high_concurrency_selection_is_aggregated_tep4_from_its_batch(): + below = table((107, 4), HC_MIN_BATCH - 1, "agg", "windowed") + assert below.param("windowed_commit_parallel", "MIN_REQUESTS") == 5 + assert below.on("wide_direct_transports") and not below.on("graph_start_qkv_prefetch") + hc = table((107, 4), HC_MIN_BATCH, "agg", "windowed") + assert hc.param("windowed_commit_parallel", "MIN_REQUESTS") == 1 + assert not hc.on("wide_direct_transports") and hc.on("graph_start_qkv_prefetch") + for key, role in (((107, 4), "gen"), ((107, 8), "agg")): + other = table(key, HC_MIN_BATCH, role, "windowed") + assert other.param("windowed_commit_parallel", "MIN_REQUESTS") == 5, (key, role) + assert not other.on("graph_start_qkv_prefetch"), (key, role) @pytest.mark.parametrize("arm", [None, B_LL, BASELINE], ids=["auto", "b-ll", "baseline"]) From 13b0b68927be290de52b8dc050cda89320ad4f8a Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Sun, 1 Nov 2026 19:19:35 -0500 Subject: [PATCH 13/15] [None][infra] Build sm_107 (Vera Rubin) into the SBSA wheel and image feat/musespark-open-agentx's aarch64 build configs and docker CUDA_ARCHS predate main's sm_107 addition, so the SBSA wheel and release image carry no Vera Rubin code, while this PR's TEP fast paths target Vera Rubin. Add 107-real to both aarch64 WHEEL_ARCHS and to the arm64 docker CUDA_ARCHS, so a side-branch release of this branch runs on Vera Rubin as well as GB300. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- docker/Makefile | 2 +- jenkins/Build.groovy | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docker/Makefile b/docker/Makefile index dcc11fc291d7..3478a87da3ce 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -38,7 +38,7 @@ DOCKER_BUILD_OPTS ?= --pull --load DOCKER_BUILD_ARGS ?= DOCKER_PROGRESS ?= auto PLATFORM ?= $(shell uname -m | grep -q 'aarch64' && echo "arm64" || echo "amd64") -CUDA_ARCHS ?= $(if $(filter arm64,$(PLATFORM)),'90-real;100-real;120-real',) +CUDA_ARCHS ?= $(if $(filter arm64,$(PLATFORM)),'90-real;100-real;103-real;107-real;120-real',) BUILD_WHEEL_OPTS ?= BUILD_WHEEL_ARGS ?= $(shell grep '^ARG BUILD_WHEEL_ARGS=' Dockerfile.multi | grep -o '=.*' | tr -d '="')$(if $(CUDA_ARCHS), --cuda_architectures $(CUDA_ARCHS))$(if $(BUILD_WHEEL_OPTS), $(BUILD_WHEEL_OPTS)) BUILD_WHEEL_SCRIPT ?= diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index b2d455e1986b..573f87a9e9fb 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -139,13 +139,13 @@ def BUILD_CONFIGS = [ (CONFIG_LINUX_AARCH64): [ (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars WARNING_IS_ERROR=ON --extra-cmake-vars NIXL_ROOT=/opt/nvidia/nvda_nixl --extra-cmake-vars MOONCAKE_ROOT=/usr/local/Mooncake --extra-cmake-vars ENABLE_BOLT_COMPATIBLE=ON", (TARNAME) : "TensorRT-LLM-GH200.tar.gz", - (WHEEL_ARCHS): "90-real;100-real;103-real;120-real", + (WHEEL_ARCHS): "90-real;100-real;103-real;107-real;120-real", (BUILD_JOBS_FOR_CONFIG): "8", // TODO: Remove after fix the build OOM issue on SBSA ], (CONFIG_LINUX_AARCH64_LLVM) : [ (WHEEL_EXTRA_ARGS) : "--extra-cmake-vars WARNING_IS_ERROR=ON --extra-cmake-vars ENABLE_BOLT_COMPATIBLE=ON -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DCMAKE_CUDA_HOST_COMPILER=clang -DCMAKE_LINKER_TYPE=LLD", (TARNAME) : "llvm-TensorRT-LLM-GH200.tar.gz", - (WHEEL_ARCHS): "90-real;100-real;103-real;120-real", + (WHEEL_ARCHS): "90-real;100-real;103-real;107-real;120-real", (BUILD_JOBS_FOR_CONFIG): "8", // TODO: Remove after fix the build OOM issue on SBSA ], ] From aa39754385d09f4823a89b836b7fa68f9eab3eb1 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Mon, 2 Nov 2026 18:49:21 -0500 Subject: [PATCH 14/15] [None][fix] Mark the prefetch module's env-var name as not a password for bandit bandit's B105 reads PRIORITIZE_FIRST_TOKEN_ENV = "TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN" as a hardcoded password (the value ends in TOKEN) and fails the Release-Check and the pre-commit workflow's bandit step. It is an environment variable's name: mark the line # nosec B105, as the other token-named constants are. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/disagg_gen_prefetch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/disagg_gen_prefetch.py b/tensorrt_llm/_torch/pyexecutor/disagg_gen_prefetch.py index 35dae095ed50..a415bb0c5832 100644 --- a/tensorrt_llm/_torch/pyexecutor/disagg_gen_prefetch.py +++ b/tensorrt_llm/_torch/pyexecutor/disagg_gen_prefetch.py @@ -44,7 +44,7 @@ DISAGG_GEN_PREFETCH_REQUESTS_ENV = "TRTLLM_DISAGG_GEN_PREFETCH_REQUESTS" # Read with the same rule as ``scheduler_v2`` so both agree on whether it is on. -PRIORITIZE_FIRST_TOKEN_ENV = "TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN" +PRIORITIZE_FIRST_TOKEN_ENV = "TLLM_DISAGG_GEN_PRIORITIZE_FIRST_TOKEN" # nosec B105 - env var name, not a password def _prioritize_first_token_enabled() -> bool: From e05bfc4592e718ecab42c1eb6ae636bc483af12e Mon Sep 17 00:00:00 2001 From: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> Date: Mon, 2 Nov 2026 18:49:21 -0500 Subject: [PATCH 15/15] [None][infra] Pass PLCScanningSetup its renamed container-scan parameters PLCScanningSetup renamed runContainerScanning to runContainerOssScanning and added runContainerMalwareScanning (main: PLC Pipeline update, #19557). This branch's L0_MergeRequest.groovy predates that and still passes runContainerScanning=false, which the job ignores: a merge request that touches an OSS-compliance file (here tensorrt_llm/usage/llm_args_golden_manifest.json) ran the container OSS scan at its default, looked for a post-merge image of the PR commit and failed [OSS-Compliance-Check]. Pass the job's current parameter names, with the values main passes, in both PLC calls. Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> --- jenkins/L0_MergeRequest.groovy | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 182d80310802..e58192c9c0a1 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -2340,7 +2340,8 @@ def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars) string(name: 'postMergeBuildNumber', value: ''), string(name: 'scanMode', value: isNightlyRelease ? 'release' : 'pre_merge'), string(name: 'runSourceCodeScanning', value: 'true'), - string(name: 'runContainerScanning', value: 'false'), + string(name: 'runContainerOssScanning', value: 'false'), + string(name: 'runContainerMalwareScanning', value: 'false'), string(name: 'runSonarQube', value: 'false'), ] def logger = new Logger(pipeline) @@ -2923,7 +2924,8 @@ def launchStages(pipeline, reuseBuild, testFilter, enableFailFast, globalVars) name: 'scanMode', value: runMode == "nightly_release" ? 'release' : 'pre_merge'), string(name: 'runSourceCodeScanning', value: 'false'), - string(name: 'runContainerScanning', value: 'true'), + string(name: 'runContainerOssScanning', value: 'true'), + string(name: 'runContainerMalwareScanning', value: 'false'), string(name: 'runSonarQube', value: 'false'), ] def logger = new Logger(pipeline)