fix: fix issuse raised by copilot

Author: windreamer

docs/en/advance/spec_decoding.md

@@ -121,8 +121,6 @@ The grammar mask is applied at two stages:
 1. **Draft model** — forked grammar matchers are used to mask each draft position serially. Each position's mask depends on the token accepted at the previous position, ensuring the draft model proposes grammatically valid tokens.
 2. **Target model verification** — position-serial grammar masking is applied to the target model's logits. After rejection sampling, only the accepted tokens are fed back to the original (un-forked) grammar matchers, keeping them in sync for the next step.
 
-For **Eagle3**, the draft vocabulary may differ from the target vocabulary, so a grammar mask cannot be directly applied to draft-vocab logits. In this case, the draft model proposes freely and grammar constraints are enforced solely on the target side via rejection sampling.
-
 ### pipeline
 
 ```python

docs/zh_cn/advance/spec_decoding.md

@@ -120,8 +120,6 @@ deepseek-ai/DeepSeek-V3 \
 1. **草稿模型** — 使用 fork 出的语法匹配器，逐位置串行施加掩码。每个位置的掩码依赖于前一位置接受的 token，确保草稿模型提出符合语法的 token。
 2. **主模型验证** — 对主模型的 logits 进行逐位置串行的语法掩码处理。拒绝采样后，仅将接受的 token 反馈给原始（未 fork 的）语法匹配器，使其为下一步保持正确的状态。
 
-对于 **Eagle3**，草稿词表可能与目标词表不同，因此无法直接对草稿词表的 logits 施加目标语法掩码。此时草稿模型自由提出 token，语法约束完全由主模型侧的拒绝采样来保证。
-
 ### pipeline
 
 ```python

lmdeploy/pytorch/spec_decode/proposers/eagle3.py

@@ -58,7 +58,7 @@ def _translate_bitmask(self, target_bitmask: torch.Tensor) -> torch.Tensor:
         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),
+        out.scatter_add_(1, draft_words.unsqueeze(0).expand(target_bitmask.size(0).to(torch.int64), -1),
                          bits_to_set)
         return out
 

lmdeploy/pytorch/spec_decode/spec_agent.py

@@ -63,7 +63,7 @@ def _expand_sampling_inputs(sampling_inputs: SamplingInputs, num_tokens: int) ->
             pass
         elif k == 'batch_size':
             v = sampling_inputs.batch_size * num_tokens
-        elif isinstance(v, (list, tuple)) and v is not None:
+        elif isinstance(v, (list, tuple)) and len(v) == sampling_inputs.batch_size:
             v = type(v)(_item for elem in v for _item in [elem] * num_tokens)
         out_dict[k] = v
 
@@ -506,8 +506,8 @@ async def _guided_spec_logits_process(
         # subsequent positions get the correct mask.
         forked = {idx: proc.fork() for idx, proc in guided_processors.items()}
 
+        guided_bitmask = guided_manager.allocate_batched_bitmap(batch_size)
         for pos in range(num_expand):
-            guided_bitmask = guided_manager.allocate_batched_bitmap(batch_size)
 
             for idx, fork_proc in forked.items():
                 guided_manager.fill_bitmap(fork_proc, guided_bitmask, idx)

tests/pytorch/spec_decode/test_guided_spec_integration.py

@@ -1104,13 +1104,22 @@ def test_get_outputs_accept_token_advances_fork(self, compiler, tokenizer_info,
         draft_token_ids, _, _ = proposer.get_outputs(
             model_outputs, model_inputs, guided_processors={0: draft_fork})
 
-        # Original matcher should be unchanged
+        # Original matcher should be unchanged, while the fork should reflect
+        # acceptance of the emitted token. Since immediate allowed-id sets may
+        # legitimately coincide for some JSON states, replay the same token on
+        # an independent fork and compare the resulting state to `draft_fork`.
+        accepted_token = int(draft_token_ids.reshape(-1)[0].item())
+        replay_fork = original.fork()
+        replay_fork.accept_token(accepted_token)
         bm_orig = guided_manager.allocate_batched_bitmap(1)
         guided_manager.fill_bitmap(original, bm_orig, 0)
         bm_fork = guided_manager.allocate_batched_bitmap(1)
         guided_manager.fill_bitmap(draft_fork, bm_fork, 0)
-        # fork has advanced by one token; original hasn't
-        assert _allowed_ids(bm_orig) != _allowed_ids(bm_fork) or True  # may coincide for JSON
+        bm_replay = guided_manager.allocate_batched_bitmap(1)
+        guided_manager.fill_bitmap(replay_fork, bm_replay, 0)
+        assert _allowed_ids(bm_fork) == _allowed_ids(bm_replay)
+        # Both original and advanced fork must remain usable for subsequent
+        # mask production even if their immediate masks happen to coincide.
 
     def test_get_outputs_multi_step_with_fork(self, compiler, tokenizer_info,
                                                 guided_manager):