feat(spec_decode): add Eagle3 guided decoding with d2t bitmask translation

- Eagle3.get_outputs() now applies grammar mask before argmax and
  accept_token after d2t mapping, matching DeepseekMTP pattern
- Add _translate_bitmask() to convert target-vocab bitmask to
  draft-vocab bitmask via scatter_add_ (vectorized, no loops)
- Remove supports_grammar_mask flag; all proposers now support it
- Fork guided processors unconditionally in spec_agent._async_model_forward
- Move session_to_cleanup handling before get_processors in forward_decode
- Bump xgrammar>=0.1.33 (fork() requirement) in all 5 runtime requirements
- Add comprehensive tests: bitmask translation, Eagle3 get_outputs,
  fork independence, multi-step draft loop

Author: windreamer

lmdeploy/pytorch/spec_decode/proposers/base.py

@@ -57,8 +57,6 @@ def draft_model_forward(
 
 class BaseSpecProposer:
 
-    supports_grammar_mask: bool = True
-
     def __init__(self, specdecode_config: SpecDecodeConfig, device: torch.device = None):
         self.specdecode_config = specdecode_config
         self.model = None

lmdeploy/pytorch/spec_decode/proposers/eagle3.py

@@ -16,16 +16,52 @@
 @SPEC_PROPOSERS.register_module(name='eagle3')
 class Eagle3(DeepseekMTP):
 
-    supports_grammar_mask: bool = False
-
     def build_model(self, empty_init: bool, target_model: torch.nn.Module = None, build_model_ctx=None):
         super().build_model(empty_init, target_model=target_model, build_model_ctx=build_model_ctx)
         self.draft_id_to_target_id = self.model.draft_id_to_target_id
+        self._init_bitmask_translate_constants()
         if not self.model.include_embed_tokens:
             logger.info('Using embed_tokens from target model.')
             del self.model.model.embed_tokens
             self.model.model.embed_tokens = target_model.get_input_embeddings()
 
+    def _init_bitmask_translate_constants(self):
+        d2t = self.draft_id_to_target_id
+        self._d2t_words = d2t // 32
+        self._d2t_bits = d2t % 32
+        draft_vocab_size = d2t.size(0)
+        draft_indices = torch.arange(draft_vocab_size, dtype=torch.int32)
+        self._draft_words = draft_indices // 32
+        self._draft_bits = draft_indices % 32
+
+    def _translate_bitmask(self, target_bitmask: torch.Tensor) -> torch.Tensor:
+        """Translate target-vocab bitmask to draft-vocab bitmask.
+
+        Args:
+            target_bitmask: [batch, ceil(target_vocab/32)] int32 bitmask
+                produced by xgr.GrammarMatcher.fill_next_token_bitmask.
+
+        Returns:
+            draft_bitmask: [batch, ceil(draft_vocab/32)] int32 bitmask
+                compatible with apply_batched_bitmap.
+        """
+        d2t_words = self._d2t_words.to(target_bitmask.device)
+        d2t_bits = self._d2t_bits.to(target_bitmask.device)
+        draft_words = self._draft_words.to(target_bitmask.device)
+        draft_bits = self._draft_bits.to(target_bitmask.device)
+
+        word_vals = target_bitmask[:, d2t_words]
+        draft_valid = ((word_vals >> d2t_bits.unsqueeze(0)) & 1).to(torch.int32)
+
+        # scatter_add_ is correct because bit positions within the same word
+        # never overlap, so addition ≡ bitwise OR.
+        bits_to_set = draft_valid << draft_bits
+        n_draft_words = (draft_valid.size(1) + 31) // 32
+        out = target_bitmask.new_zeros(target_bitmask.size(0), n_draft_words)
+        out.scatter_add_(1, draft_words.unsqueeze(0).expand(target_bitmask.size(0), -1),
+                         bits_to_set)
+        return out
+
     def get_target_hidden_size(self, model_config: ModelConfig):
         """Get target hidden size."""
         hf_config = self.specdecode_config.model_config.hf_config
@@ -35,7 +71,8 @@ def get_target_hidden_size(self, model_config: ModelConfig):
     def get_outputs(self,
                     model_outputs: dict[str, torch.Tensor],
                     model_inputs: ModelInputs,
-                    extra_inputs: ExtraInputs = None):
+                    extra_inputs: ExtraInputs = None,
+                    guided_processors: dict | None = None):
         """Get outputs."""
         hidden_states = model_outputs['hidden_states']
         hidden_states_prenorm = model_outputs['hidden_states_prenorm']
@@ -52,7 +89,20 @@ def get_outputs(self,
 
         logits = self.get_logits(hidden_states)[0]
 
+        if guided_processors and self.guided_decoding_manager is not None:
+            guided_manager = self.guided_decoding_manager
+            guided_bitmask = guided_manager.allocate_batched_bitmap(logits.size(0))
+            for idx, processor in guided_processors.items():
+                guided_manager.fill_bitmap(processor, guided_bitmask, idx)
+            draft_bitmask = self._translate_bitmask(guided_bitmask)
+            guided_manager.apply_batched_bitmap(logits, draft_bitmask)
+
         draft_token_ids = logits.argmax(dim=-1, keepdim=True)
         draft_token_ids = self.draft_id_to_target_id[draft_token_ids]
 
+        if guided_processors and self.guided_decoding_manager is not None:
+            guided_manager = self.guided_decoding_manager
+            for idx, processor in guided_processors.items():
+                guided_manager.accept_token(processor, draft_token_ids[idx, 0].item())
+
         return draft_token_ids, model_metas, hidden_states_prenorm

lmdeploy/pytorch/spec_decode/spec_agent.py

@@ -368,9 +368,15 @@ def __compute_logprobs(raw_logprobs: torch.Tensor, token_ids: torch.LongTensor,
 
         guided_processors = {}
         guided_manager = self.guided_decoding_manager
-        if guided_manager and sampling_inputs.session_ctx is not None:
-            guided_processors = guided_manager.get_processors(
-                sampling_inputs.session_ctx, sampling_inputs.response_formats)
+        if guided_manager:
+            session_to_cleanup = sampling_inputs.session_to_cleanup
+            if session_to_cleanup is not None:
+                for session_id in session_to_cleanup:
+                    guided_manager.remove_processor(session_id)
+
+            if sampling_inputs.session_ctx is not None:
+                guided_processors = guided_manager.get_processors(
+                    sampling_inputs.session_ctx, sampling_inputs.response_formats)
 
         if model_inputs.is_decoding:
             if guided_processors:
@@ -539,11 +545,10 @@ async def _async_model_forward(self, inputs: ModelInputs, extra_inputs: ARSpecEx
             # create dummy draft tokens
             output_draft_ids = inputs.input_ids.new_zeros(1, self.num_spec_tokens)
         else:
-            # Fork guided processors when the proposer supports grammar masking.
+            # Fork guided processors for draft model.
             draft_guided_processors = None
             guided_manager = self.guided_decoding_manager
-            if (guided_manager and sampling_inputs.session_ctx is not None
-                    and self.proposer.supports_grammar_mask):
+            if guided_manager and sampling_inputs.session_ctx is not None:
                 orig_processors = guided_manager.get_processors(
                     sampling_inputs.session_ctx, sampling_inputs.response_formats)
                 draft_guided_processors = {idx: proc.fork()

requirements/runtime_ascend.txt

@@ -27,4 +27,4 @@ torch-npu>=2.3.1,<2.10.0
 torchvision>=0.18.1,<0.25.0
 transformers
 uvicorn
-xgrammar
+xgrammar>=0.1.33

requirements/runtime_camb.txt

@@ -26,4 +26,4 @@ torch<=2.6.0,>=2.4.0
 torchvision<=0.21.0,>=0.15.0
 transformers
 uvicorn
-xgrammar
+xgrammar>=0.1.33

requirements/runtime_cuda.txt

@@ -29,4 +29,4 @@ torchvision<=0.25.0,>=0.15.0
 transformers>=4.52.0
 triton<=3.6.0,>=3.0.0; sys_platform == "linux" and "aarch64" not in platform_machine and "arm" not in platform_machine
 uvicorn
-xgrammar
+xgrammar>=0.1.33

requirements/runtime_maca.txt

@@ -27,4 +27,4 @@ torchvision<=0.21.0,>=0.15.0
 transformers
 triton>=2.1.0; sys_platform == "linux"
 uvicorn
-xgrammar
+xgrammar>=0.1.33

requirements/runtime_rocm.txt

@@ -21,4 +21,4 @@ shortuuid
 tiktoken
 transformers
 uvicorn
-xgrammar
+xgrammar>=0.1.33