atom(pd): every P/D producer reports its first token from the forward that samples it

atom/compass/design/09_fitting_and_law_selection.md

@@ -738,8 +738,9 @@ refused, unless a peer prefills (below).
 **Assumed treatment: eager decode beside prefill** (owner ruling on
 [#2660](https://github.com/jgong5/ATOM/issues/2660#issuecomment-6573151549), 2026-11-02).
 A DP group replays request graphs only on a step every rank runs alike, so a rank that
-decodes while a peer prefills runs eagerly. A 1P1D prefill side runs such a decode for
-each request its deferred output holds (#2655). The fit cuts a step whose ranks fall to
+decodes while a peer prefills runs eagerly. A 1P1D prefill side ran such a decode for
+each request its deferred output held, until #2655 had every producer report its first
+token from the forward that samples it. The fit cuts a step whose ranks fall to
 different vectors, so no law fits an `eager` vector for it. A step linear law with no
 `eager` vector and a replay step grid price it at 0 s under this treatment, counted by
 step like the dummy's: the rank waits on its prefilling peers inside the forward, and the

atom/compass/runner/step_output.py

@@ -9,7 +9,7 @@
 
 **Tokens are one step late.** `tokenIDProcessor.is_deferred_out` is true
 whenever the pipeline is a single stage, which is the ordinary case, except on
-a P/D producer whose model keeps recurrent state. The real
+a P/D producer. The real
 runner starts an asynchronous copy of this step's sampled ids to the host and
 reports the copy the *previous* step started, so the reply carries the previous
 batch's request ids and the previous batch's tokens with the deferred flag set.

atom/model_engine/model_runner.py

@@ -224,17 +224,15 @@ def __init__(
         # The scheduler never speculates for a P/D producer, though the runner
         # keeps its drafter loaded: every decode it runs is q=1.
         self.is_pd_producer = _kv_config_has_producer(kv_cfg)
-        # P/D hands off prompt-end state plus the first sampled token. Deferred
-        # output makes the producer decode once more to surface that token, so
-        # the consumer applies T0 to a state that already has it. Only a
-        # recurrent state notices -- a paged write lands at T0's own position,
-        # so repeating it is idempotent.
-        hands_off_recurrent_state = (
-            self.is_pd_producer and runner.attn_family.has_recurrent_state
-        )
-        self.is_deferred_out = (
-            not self.is_pipeline_parallel and not hands_off_recurrent_state
-        )
+        # P/D hands off prompt-end state plus the first sampled token (T0), and
+        # a producer releases a request only once T0 reaches the host. Deferred
+        # output surfaces T0 at the rank's next forward that samples: another
+        # request's output prefill, or one more decode of the held request when
+        # nothing else is scheduled. A forward of middle chunks samples nothing,
+        # so the request waits while the rank prefills others (#2655: a median
+        # 14 s). That extra decode also makes a recurrent state apply T0 twice.
+        # So every producer reports T0 from the forward that samples it.
+        self.is_deferred_out = not self.is_pipeline_parallel and not self.is_pd_producer
 
         self.runner = runner
         device = runner.device
@@ -2577,8 +2575,8 @@ def _verifies_drafts(self, batch: ScheduledBatch, is_prefill: bool) -> bool:
         """Whether this forward scores drafts with the rejection sampler.
 
         A P/D producer keeps the drafter only to write its context KV; the
-        scheduler never pads drafts for it, so its decodes (the step that
-        surfaces a deferred T0) are q=1 and have nothing to verify.
+        scheduler never pads drafts for it, so its decodes are q=1 and have
+        nothing to verify.
         """
         return (
             not is_prefill

atom/model_engine/scheduler.py

@@ -703,9 +703,8 @@ def __init__(
 
         # Speculative decoding. A P/D producer is excluded for the same reason
         # `PrefillScheduler` is: it prefills, samples T0 and hands the request
-        # off, so it never proposes. That holds for every producer, so this
-        # does not follow the runner's deferred-output decision, which
-        # additionally asks whether the model keeps recurrent state.
+        # off, so it never proposes. That holds for every producer, as does
+        # the runner's giving up deferred output.
         # Leaving speculation on here would still pad every sequence with
         # `mtp_k` placeholders and size its decode window to `mtp_k + 1`, and
         # the padding holds `num_tokens` `mtp_k` short of `max_tokens` so the

tests/compass/test_capture_real_model.py

@@ -200,7 +200,7 @@
 ROW_PARALLEL = "atom/model_ops/layernorm.py:944 in _forward_fused_allreduce"
 VOCAB_EMBEDDING = "atom/model_ops/embed_head.py:178 in forward"
 VOCAB_LM_HEAD = "atom/model_ops/embed_head.py:269 in forward"
-SAMPLER = "atom/model_engine/model_runner.py:3389 in postprocess"
+SAMPLER = "atom/model_engine/model_runner.py:3387 in postprocess"
 
 # Where a free symbol stops being free, each as the value it takes and the
 # innermost ATOM frames it takes it through. Three, and in an order: a symbol
@@ -505,10 +505,10 @@
     "atom/model_ops/attentions/backends.py:671 in _build_mrope_decode_positions": 1,
     "atom/model_ops/attentions/backends.py:714 in prepare_cu_seqlens_q": 1,
     "atom/model_ops/attentions/backends.py:715 in prepare_cu_seqlens_q": 3,
-    "atom/model_engine/model_runner.py:522 in prepare_input_ids": 3,
-    "atom/model_engine/model_runner.py:608 in _stage_input_ids": 1,
-    "atom/model_engine/model_runner.py:611 in _stage_input_ids": 1,
-    "atom/model_engine/model_runner.py:2698 in prepare_sample": 1,
+    "atom/model_engine/model_runner.py:520 in prepare_input_ids": 3,
+    "atom/model_engine/model_runner.py:606 in _stage_input_ids": 1,
+    "atom/model_engine/model_runner.py:609 in _stage_input_ids": 1,
+    "atom/model_engine/model_runner.py:2696 in prepare_sample": 1,
     "atom/model_ops/sampler.py:48 in _apply_greedy_tokens": 1,
     "atom/utils/forward_context.py:318 in decide": 1,
     "atom/utils/h2d.py:51 in _count": WIDTH_COUNTED_COPIES,

tests/compass/test_runner_step_semantics.py

@@ -359,10 +359,10 @@ def test_a_pipeline_stage_reports_the_step_it_ran(run):
 def test_a_single_stage_defers_exactly_when_atoms_token_processor_does(
     run, monkeypatch, deferred
 ):
-    """ATOM's `tokenIDProcessor` clears `is_deferred_out` under PP and on a P/D
-    producer whose model keeps recurrent state (a GDN hybrid), and the stream
-    the first predicted step builds takes that answer. So a single-stage
-    producer of such a model reports the decode it ran, not the one before."""
+    """ATOM's `tokenIDProcessor` clears `is_deferred_out` under PP and on every
+    P/D producer, and the stream the first predicted step builds takes that
+    answer. So a single-stage producer reports the step it ran, not the one
+    before."""
     from atom.model_engine import llm_engine
 
     monkeypatch.setattr(llm_engine, "_load_tokenizer", lambda *_: None)

tests/test_spec_deferred_status_queue.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import pytest
+import torch
 
 # On the module, not on "aiter" -- see tests/test_dspark.py for why a name-only
 # guard passes on the non-GPU runner and then aborts collection.
@@ -179,33 +180,46 @@ def test_a_consumer_keeps_deferred_output():
 
 
 @pytest.mark.parametrize(
-    ("family", "expected"),
+    "family",
     [
-        (Family.KIMI_MLA, False),
-        (Family.QSA_GDN, False),
-        (Family.GDN, False),
-        (Family.V4, True),
-        (Family.MLA, True),
-        (Family.MHA, True),
-        (Family.CSA2, False),
+        Family.KIMI_MLA,
+        Family.QSA_GDN,
+        Family.GDN,
+        Family.V4,
+        Family.MLA,
+        Family.MHA,
+        Family.CSA2,
     ],
 )
-def test_only_recurrent_state_producers_give_up_deferred_output(family, expected):
-    """A paged write lands at T0's own position, so the consumer repeating it
-    changes nothing -- V4 and the other paged models stay on the path they
-    used before P/D handoff existed. CSA2 (DeepSeek-V4.1) hands off a STATE
-    slot whose cursor, compressor rings and Engram history advance per token,
-    so like the recurrent families its producer emits T0 from the prefill."""
+def test_every_producer_gives_up_deferred_output(family):
+    """A producer releases a request once T0 reaches the host. Deferred, T0
+    waits for the rank's next forward that samples, which middle chunks of
+    other requests push back by seconds (#2655) -- paged (V4) and recurrent
+    (GDN, CSA2) state alike."""
     processor = _processor_for(kv_config={"kv_role": "kv_producer"}, family=family)
-    assert processor.is_deferred_out is expected
+    assert processor.is_deferred_out is False
+
+
+def test_a_producer_reports_t0_from_the_forward_that_samples_it():
+    """The output prefill's own forward returns its sampled ids, with the
+    flag that tells the scheduler they are not one step late."""
+    processor = _processor_for(kv_config={"kv_role": "kv_producer"}, family=Family.V4)
+    batch = SimpleNamespace(req_ids=[5, 7])
+
+    ids, logprobs = processor.prepare_sampled_ids(
+        batch, torch.tensor([42, 43], dtype=torch.int32), sync_event=None
+    )
+
+    assert ids == {5: (42,), 7: (43,), -1: 0}
+    assert logprobs is None
 
 
 @pytest.mark.parametrize(
     ("kv_role", "expected"), [("kv_producer", True), ("kv_consumer", False)]
 )
 def test_the_processor_records_the_producer_role_whatever_the_state(kv_role, expected):
-    """The runner reads it to size its DP dummy and to scope its graph lookup,
-    and a paged producer -- which keeps deferred output -- is still one."""
+    """The runner reads it to size its DP dummy, to scope its graph lookup and
+    to skip verifying drafts, whatever state the model hands off."""
     processor = _processor_for(kv_config={"kv_role": kv_role}, family=Family.V4)
     assert processor.is_pd_producer is expected
 
@@ -223,7 +237,7 @@ def _decode_batch(*, num_spec_step: int) -> ScheduledBatch:
     return batch
 
 
-def _recurrent_producer_processor() -> tokenIDProcessor:
+def _producer_processor() -> tokenIDProcessor:
     """Non-deferred, and `use_spec` still on: the runner keeps the drafter."""
     processor = _processor()
     processor.is_deferred_out = False
@@ -233,7 +247,7 @@ def _recurrent_producer_processor() -> tokenIDProcessor:
 
 
 def test_a_non_deferred_decode_stages_the_scheduled_ids_as_they_are():
-    processor = _recurrent_producer_processor()
+    processor = _producer_processor()
 
     tokenIDProcessor.prepare_input_ids(processor, _decode_batch(num_spec_step=0), 1)
 
@@ -246,7 +260,7 @@ def test_the_non_deferred_branch_refuses_a_batch_carrying_drafts():
     never speculates. Staging them here would size rows by `max_seqlen_q`,
     which under DP is the group's maximum for a block drafter, not this
     batch's row width."""
-    processor = _recurrent_producer_processor()
+    processor = _producer_processor()
 
     with pytest.raises(AssertionError, match="num_spec_step=3"):
         tokenIDProcessor.prepare_input_ids(processor, _decode_batch(num_spec_step=3), 4)
@@ -310,7 +324,7 @@ def test_only_a_producer_with_a_drafter_may_miss_a_recording():
     ("is_pd_producer", "drafter", "is_prefill", "is_dummy_run", "expected"),
     [
         (False, True, False, False, True),
-        # The producer's decode surfaces a deferred T0 at q=1: no drafts, so a
+        # A producer decodes at q=1: it never speculates, so a
         # rejection pass would select over zero target rows.
         (True, True, False, False, False),
         (False, True, True, False, False),