[EtsyML] Native TRIE guided decoding (per-beam, host-staging)

.dockerignore

@@ -12,3 +12,8 @@ examples/**/*.onnx
 examples/**/*.safetensors
 examples/**/c-model
 examples/models/core/gpt/gpt*
+
+cpp/build
+.git
+*.whl
+trie_decoding

cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2024-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -21,6 +21,10 @@
 #include "tensorrt_llm/runtime/bufferManager.h"
 #include "tensorrt_llm/runtime/iTensor.h"
 
+#include <memory>
+#include <mutex>
+#include <unordered_map>
+
 namespace xgrammar
 {
 class GrammarMatcher;
@@ -31,6 +35,9 @@ namespace tensorrt_llm::batch_manager
 {
 class DecoderInputBuffers;
 
+// Forward declaration for trie-based constraint handling (flat CSR representation)
+class GuidedTrie;
+
 class GuidedDecoder
 {
 public:
@@ -39,29 +46,54 @@ class GuidedDecoder
     using BitmaskT = uint32_t;
 
     GuidedDecoder(executor::GuidedDecodingConfig const& guidedDecodingConfig, SizeType32 maxNumSequences,
-        SizeType32 vocabSizePadded, nvinfer1::DataType logitsDtype, runtime::BufferManager const& runtimeBufferManager);
+        SizeType32 maxBeamWidth, SizeType32 vocabSizePadded, nvinfer1::DataType logitsDtype,
+        runtime::BufferManager const& runtimeBufferManager);
     void build(ScheduledRequests const& scheduledRequests);
     void execute(DecoderInputBuffers const& decoderInputBuffers, runtime::BufferManager const& runtimeBufferManager);
 
+    /// @brief TRIE constraints are applied per beam; at beam width > 1 the caller must refresh
+    /// each request's beam token histories (gatherTree) before build(), because the raw slot
+    /// buffers go stale when beam search switches hypothesis parents.
+    [[nodiscard]] bool needsGatheredBeamTokens() const
+    {
+        return mGuidedDecodingBackend == executor::GuidedDecodingConfig::GuidedDecodingBackend::kTRIE;
+    }
+
 private:
     executor::GuidedDecodingConfig::GuidedDecodingBackend mGuidedDecodingBackend;
     std::vector<std::shared_ptr<xgrammar::GrammarMatcher>> mXGrammarMatchers;
     std::shared_ptr<xgrammar::GrammarCompiler> mXGrammarCompiler;
 
+    // Trie-based constraint handling (flat CSR trie; nodes are int32 indices, root = 0).
+    // No persistent per-request node state: each step re-walks every beam's generated tokens
+    // (bounded by the trie depth) so beam reordering cannot desynchronize the constraint.
+    std::shared_ptr<GuidedTrie> mTrie;
+    std::unordered_map<int32_t, std::vector<BitmaskT>> mTrieBitmaskCache;
+    std::vector<BitmaskT> mEosOnlyBitmask; // fallback for off-trie beams: allow only EOS
+    std::mutex mTrieCacheMutex;
+
     SizeType32 mMaxNumSequences;
+    SizeType32 mMaxBeamWidth;
     SizeType32 mVocabSizePadded;
     SizeType32 mBitmaskSize; // CeilDiv(vocabSizePadded, 32)
     nvinfer1::DataType mLogitsDtype;
 
-    TensorPtr mLogitsBitmask;           // [mMaxNumRequests, mBitmaskSize]
-    TensorPtr mLogitsBitmaskHost;       // [mMaxNumRequests, mBitmaskSize]
-    TensorPtr mLogitsBitmaskPtrVec;     // [mMaxNumRequests], pointers to the logitsBitmask in a batch
-    TensorPtr mLogitsBitmaskPtrVecHost; // [mMaxNumRequests]
-    TensorPtr mLogitsPtrVec;            // [mMaxNumRequests], pointers to the logits in a batch
-    TensorPtr mLogitsPtrVecHost;        // [mMaxNumRequests]
+    // Bitmask row layout: one row per (seqSlot, beam): row = seqSlot * mMaxBeamWidth + beam.
+    // The xgrammar backend keeps its per-request design and uses beam 0's row.
+    TensorPtr mLogitsBitmask;           // [mMaxNumSequences * mMaxBeamWidth, mBitmaskSize]
+    TensorPtr mLogitsBitmaskHost;       // [mMaxNumSequences * mMaxBeamWidth, mBitmaskSize]
+    TensorPtr mLogitsBitmaskPtrVec;     // [mMaxNumSequences * mMaxBeamWidth]
+    TensorPtr mLogitsBitmaskPtrVecHost; // [mMaxNumSequences * mMaxBeamWidth]
+    TensorPtr mLogitsPtrVec;            // [mMaxNumSequences * mMaxBeamWidth]
+    TensorPtr mLogitsPtrVecHost;        // [mMaxNumSequences * mMaxBeamWidth]
 
     // BufferManager with a dedicated stream for async copy of buffers for guided decoding.
     runtime::BufferManager mCopyBufferManager;
+
+    // Helper methods for trie
+    void initTrieBuffers(runtime::BufferManager const& runtimeBufferManager);
+    std::vector<BitmaskT> const& getTrieBitmask(int32_t node);
+    void createTrieBitmask(int32_t node, std::vector<BitmaskT>& bitmask);
 };
 
 } // namespace tensorrt_llm::batch_manager

cpp/include/tensorrt_llm/executor/executor.h

@@ -1431,12 +1431,15 @@ class GuidedDecodingConfig
         kXGRAMMAR = 0,
         /// @brief Enable guided decoding with LLGuidance backend.
         kLLGUIDANCE = 1,
+        /// @brief Enable guided decoding with Trie-based constraint backend.
+        kTRIE = 2,
     };
 
     explicit GuidedDecodingConfig(GuidedDecodingBackend backend,
         std::optional<std::vector<std::string>> encodedVocab = std::nullopt,
         std::optional<std::string> tokenizerStr = std::nullopt,
-        std::optional<std::vector<TokenIdType>> stopTokenIds = std::nullopt);
+        std::optional<std::vector<TokenIdType>> stopTokenIds = std::nullopt,
+        std::optional<std::string> triePath = std::nullopt);
 
     bool operator==(GuidedDecodingConfig const& other) const;
 
@@ -1452,6 +1455,9 @@ class GuidedDecodingConfig
     void setStopTokenIds(std::vector<TokenIdType> const& stopTokenIds);
     [[nodiscard]] std::optional<std::vector<TokenIdType>> getStopTokenIds() const;
 
+    void setTriePath(std::string const& triePath);
+    [[nodiscard]] std::optional<std::string> getTriePath() const;
+
     void validate() const;
 
 private:
@@ -1472,6 +1478,8 @@ class GuidedDecodingConfig
     std::optional<std::string> mTokenizerStr;
     /// @brief Stop token ids. If not provided, it can be automatically detected.
     std::optional<std::vector<TokenIdType>> mStopTokenIds;
+    /// @brief Path to binary trie file for TRIE backend.
+    std::optional<std::string> mTriePath;
 };
 
 class LogitsPostProcessorConfig