refactor(grammar): move mask dispatch to stream seam

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: wanglining97

docs/backend/xgrammar_design.md

@@ -55,7 +55,7 @@ above speaks generic logits-processor vocabulary.
         ▼
   stream / executor / scheduler  ── NO
    GenerateStream (gates on prepare()), NormalExecutor / MtpExecutor
-   (call batch_logits_processor::*), FIFOScheduler (generic readiness skip),
+   (call stream_logits_processor::*), FIFOScheduler (generic readiness skip),
    BatchDecodeScheduler (only the generic hasStructuredOutputRequest() predicate)
         │ wired once at startup by
         ▼
@@ -71,11 +71,11 @@ Two seams cross these layers without leaking concrete types:
    generic factory, breaking the otherwise-cyclic build dependency
    (`models:logits_processor` → xgrammar → `engine_base/stream` →
    `models:logits_processor`).
-2. **Batch-dispatch seam** — `batch_logits_processor::*` + `SpecLogitsVerifyRunner`
-   let executors drive whole-batch masks (verify / draft / decode). The dispatch
-   is the only code that walks `StreamGroups`; it hands the runner a flat list of
-   `SpecLogitsProcessor` participants, so the model layer and executors never see
-   `StreamGroups` or a concrete backend.
+2. **Stream-dispatch seam** — `stream_logits_processor::*` + `SpecLogitsVerifyRunner`
+   lets stream/executor preparation code drive whole-batch masks (verify / draft /
+   decode). The dispatch is the only code that walks `StreamGroups`; it hands the
+   runner a flat list of `SpecLogitsProcessor` participants, so the model layer and
+   concrete executors never see a constraint backend.
 
 Data (not behavior) does cross the boundary in two benign forms: the request
 config fields on `GenerateConfig` (`json_schema` / `regex` / `ebnf` /
@@ -116,10 +116,10 @@ configuration contracts, not engine-type coupling.
    that exposes only per-stream primitives (`tryAcceptAndFillBitmask` for verify,
    `fillDecodeBitmask` / `fillDraftBitmask` for decode/draft, plus a generic
    `applyBitmask`). For all three batch passes (verify, draft, decode),
-   `BatchLogitsProcessorDispatch` collects the active processors from
+   `StreamLogitsProcessorDispatch` collects the active processors from
    `StreamGroups` into a flat list and `SpecLogitsVerifyRunner` owns the batch
    assembly (allocate, per-stream fill, merge, apply). The model layer never sees
-   `StreamGroups`; callers name only the generic `batch_logits_processor`
+   `StreamGroups`; callers name only the generic `stream_logits_processor`
    namespace.
 
 5. Fail closed rather than silently generating unconstrained output.
@@ -251,7 +251,7 @@ calls `notifyCommit`.
 ### Normal Decode Batch Mask
 
 `NormalExecutor` calls
-`batch_logits_processor::batchApplyDecodeMaskFromStreams` immediately before
+`stream_logits_processor::applyDecodeMaskFromStreams` immediately before
 sampling, which collects participants and runs `SpecLogitsVerifyRunner::runDecode`
 (one row per stream). This path collapses per-stream mask launches into one batch
 apply. If the batch mask ran, `SamplerInputs::is_decode_mask_batched` is set so the
@@ -268,19 +268,21 @@ model layer):
 
 - draft-side masking (`runDraft` + `fillDraftBitmask`) keeps proposed tokens in
   the legal set when possible: each processor DFS-walks its proposed chain, fills
-  its single row, then rolls back so observable state is unchanged;
+  its single row, then rolls back so observable state is unchanged. MTP exposes
+  this through `MtpDraftMaskState`, so the executor does not manage CPU draft-chain
+  row alignment itself;
 - verify-side masking (`run` + `tryAcceptAndFillBitmask`) fills each stream's
   `[propose_step + 1, words]` row block, merges processors for the same stream by
   bitwise AND, and applies the merged bitmask through the generic `applyBitmask()`
-  hook.
+  hook. This is applied while `MtpBatchStreamProcessor` prepares the target
+  sampler input.
 
 The actual matcher state remains unchanged during those batch passes. After the
 speculative sampler decides the accepted suffix, the normal stream commit path
 updates every attached logits processor exactly once.
 
-`MtpExecutor` and `MtpBatchStreamProcessor` therefore know only about
-`batch_logits_processor`, `SpecLogitsVerifyRunner`, and `SpecLogitsProcessor`,
-not `GrammarLogitsProcessor`.
+`MtpExecutor` and `MtpBatchStreamProcessor` therefore know only about the
+`stream_logits_processor` seam, not `GrammarLogitsProcessor`.
 
 ### PD Separation
 
@@ -340,7 +342,7 @@ resolve the hook's definition). This keeps the composition root explicit (no
 
 `NormalExecutor`, `MtpExecutor`, and `MtpBatchStreamProcessor` should not know
 about grammar. The current design satisfies that by routing all batch mask work
-(verify, draft, decode) through `batch_logits_processor`: the stream dispatch
+(verify, draft, decode) through `stream_logits_processor`: the stream dispatch
 collects `SpecLogitsProcessor` participants into a flat list and delegates the
 actual assembly to `SpecLogitsVerifyRunner`, so no concrete processor owns the
 whole batch and the model layer never sees `StreamGroups`. The names are

rtp_llm/cpp/engine_base/schedulers/FIFOScheduler.cc

@@ -247,10 +247,10 @@ void FIFOScheduler::evaluateWaitingStreams(list<GenerateStreamPtr>& waiting_stre
         }
 
         // Skip streams that are still waiting for async preparation (e.g.
-        // grammar compile). They stay in waiting_streams_ until prepare()
-        // completes, which prevents them from being admitted into a batch
-        // that may contain decode streams (avoiding mixed prefill+decode).
-        if (stream->isPreparationPending()) {
+        // grammar compile) AND haven't been externally marked CanRun (e.g.
+        // by DecodeRpcServer for PD decode streams). Externally-marked streams
+        // get their preparation polled through moveToNext in Phase 2.
+        if (stream->isPreparationPending() && !stream->hasEvent(StreamEvents::CanRun)) {
             it++;
             continue;
         }

rtp_llm/cpp/engine_base/stream/BatchLogitsProcessorDispatch.cc

@@ -453,6 +453,17 @@ class GenerateStream: public std::enable_shared_from_this<GenerateStream> {
         logits_processor_list_.push_back(std::move(processor));
     }
 
+    // Post-construction init: set stream back-references on processors that
+    // need it (e.g. GrammarLogitsProcessor for PD token replay). Must be called
+    // after make_shared so shared_from_this() is valid.
+    void initProcessorStreamRefs() {
+        for (auto& p : logits_processor_list_) {
+            if (p) {
+                p->setStream(shared_from_this());
+            }
+        }
+    }
+
     // Generic processor lookup by concrete type. Returns the first attached
     // processor whose dynamic type is T (or a subclass), or nullptr.
     // Use in lieu of grammar-specific finders so callers depend on the