Harden dynamic inference request lifecycle handling

[lmcafee-nvidia]

What does this PR do?

This PR hardens dynamic inference across request admission, checkpoints, pause and resume, recompute, completion, and result construction. It preserves request identity and state, bounds prompt memory residency, aligns token metadata, and returns failures and completed text through consistent ownership boundaries.

It combines the active request-lifecycle split review records:

• Prompt serialization and checkpoint state: fork/47, fork/48, and fork/49.
• Completion, reset, admission, and failure ownership: fork/51, fork/52, fork/53, and fork/54.
• Pause, recompute, stop, and score correctness: fork/55, fork/56, fork/57, and fork/58.
• Pairwise request-lifecycle and memory-residency coverage: fork/59.

Example failure scenario(s)

Scenario 1 (Hidden prompt remains in a result)

• Request A has 4 prompt tokens.
• A does not ask for prompt tokens in its result.
• One prompt view is hidden but another still exposes the same tokens.
• The fix omits both prompt views unless the caller asks for them.

Scenario 2 (Serialization damages the live request)

• Request A has 4 live prompt tokens.
• Serialization fails while building A's result.
• A is left without its prompt state.
• The fix restores the live state even when serialization fails.

Scenario 3 (Checkpoint loses request settings)

• Request A may generate 5 tokens and has custom settings.
• A generates 2 tokens and checkpoints.
• A resumes without part of its settings, status, timing, or history.
• The fix preserves independent copies of that state and the remaining 3-token budget.

Scenario 4 (Checkpoint history retains GPU memory)

• Request A starts with 4 prompt tokens.
• A checkpoints 3 times and creates progressively larger prompt tensors.
• Old intermediate prompts remain allocated on the GPU.
• The fix moves inactive intermediate prompts to CPU while retaining the two active prompt allocations.

Scenario 5 (Checkpoint boundary changes decoded text)

• Request A generates tokens 10 and 11 before a checkpoint.
• A generates token 12 after the checkpoint.
• Decoding each segment separately can differ from decoding tokens 10, 11, and 12 together.
• The fix merges tokens once and decodes the complete stream at the caller boundary.

Scenario 6 (Reset strands a coordinator waiter)

• Request A finishes and the engine is drained.
• A coordinator task still waits on the engine condition.
• Reset replaces that condition with a new object.
• The fix preserves the long-lived coordinator objects so later notifications reach the waiter.

Scenario 7 (Reset mutates suspended storage)

• The engine enters its suspended state and releases inference storage.
• Reset begins mutating the unavailable context.
• The fix rejects reset before any mutation occurs.

Scenario 8 (Omitted sampling settings fail admission)

• Request A omits sampling settings.
• Admission reads a setting before creating defaults.
• A fails instead of receiving standard settings.
• The fix creates a fresh default settings object for A.

Scenario 9 (Duplicate ID mutates state)

• Request A is active with ID 7.
• Request B also arrives with ID 7.
• B starts admission work against A's existing entry.
• The fix rejects B before either request is mutated.

Scenario 10 (An all-invalid call returns nothing)

• Call A submits 2 invalid requests.
• Both requests fail during admission and need no model step.
• The old loop exits with 0 results.
• The fix returns both failures immediately.

Scenario 11 (One caller consumes another caller's failure)

• Call A owns failed request 9.
• Call B submits requests 0 and 1.
• Global cleanup for B also removes request 9.
• The fix limits each call to the failures from its own submitted futures.

Scenario 12 (A coordinator failure remains after reply)

• The coordinator submits invalid request 42.
• The engine sends request 42's failure reply immediately.
• No valid request remains to trigger another bookkeeping step.
• The fix removes request 42 at the end of its admission batch without replying twice.

Scenario 13 (Paused ID has the wrong shape)

• Requests A and B are active.
• B needs another cache block and pauses.
• The old lookup returns a nested value such as [[2]] instead of ID 2.
• The fix returns B's flat, stable request ID.

Scenario 14 (Resume removes the wrong pause)

• Requests A and B are paused.
• Memory pressure changes which request can resume.
• Removing a list suffix clears the wrong paused entry.
• The fix removes only the request ID that actually resumed.

Scenario 15 (Recompute changes request order)

• Requests A and B are active while C and D are paused.
• The resident rows are C, D, A, and B.
• An unordered replay assigns random draws and metadata to different requests.
• The fix replays A, B, D, and C deterministically.

Scenario 16 (Partial prompt scores are counted twice)

• Request A has a 6-token prompt.
• A processes 3 prompt tokens and suspends before generation.
• Resume restarts all 6 prompt tokens but retains the first 3 scores.
• The fix clears those partial scores before recomputation.

Scenario 17 (Completed prompt scores are erased)

• Request A finishes prefill and generates token 7.
• A suspends with a nonzero finished-chunk count.
• A broad partial-prefill reset erases its completed prompt scores.
• The fix resets scores only when no generated token exists.

Scenario 18 (A stop sequence crosses a checkpoint)

• Request A stops on tokens 20 and 21.
• A generates token 20 and checkpoints.
• The next segment generates token 21 but checks only its local segment.
• The fix detects the complete stop sequence across both segments.

Scenario 19 (Speculation passes the first stop)

• Request A stops on tokens 30 and 31.
• One speculative step returns tokens 30, 31, and 32.
• Token 32 and its metadata remain after the first stop.
• The fix keeps the earliest stop and trims every later aligned value.

Scenario 20 (A stripped step has no score to read)

• Request A skips prompt log probabilities.
• A completes a stop sequence at a checkpoint boundary.
• Removing the stop sequence leaves no generated token or score in that step.
• The fix appends nothing unless both a token and a score survive.

Scenario 21 (Partial prefill scores the wrong token)

• Request A has prompt tokens 10, 11, 12, and 13.
• A chunk processes tokens 10 and 11, then the sampler proposes token 99.
• The engine discards token 99 because known prompt token 12 comes next, but it saves token 99's score.
• The fix saves the score for token 12.

Scenario 22 (Skipped prompt scores leave a hidden row)

• Request A disables prompt log probabilities.
• A partial chunk still computes a provisional top-N row.
• The old path retains that row even though no prompt score was requested.
• The fix omits the row together with the skipped prompt scores.

How was this bug discovered?

The campaign compared uninterrupted requests with requests interrupted at lifecycle boundaries and checked exact tokens, text, scores, state, identities, memory residency, and completion ownership.

Split PR	Observable clue and oracle
fork/47	Payload inspection exposed a second prompt tensor, and an injected serialization failure showed that the live request was not restored.
fork/48	A checkpoint-state comparison exposed missing settings, an incorrect token budget, lost status and timing, and shared mutable history.
fork/49	CUDA storage identities showed that every cumulative checkpoint prompt remained reachable.
fork/51	Token IDs matched while text differed under a tokenizer whose segmented and complete-stream decoding are intentionally different.
fork/52	Exact runtime-state checks found replaced coordinator primitives and context mutation during an invalid suspended reset.
fork/53	Direct admission checks exposed the missing default and showed duplicate IDs reaching mutable work.
fork/54	An all-invalid batch returned no failures, interleaved callers crossed ownership boundaries, and an immediate coordinator failure remained resident.
fork/55	Memory pressure produced an unhashable paused ID, and exact event checks found the wrong pause entry after resume.
fork/56	Uninterrupted and resumed runs assigned outputs differently, duplicated partial scores, or erased completed scores.
fork/57	Exact token, score, event, and timing comparisons exposed missed or misaligned stop boundaries.
fork/58	A deliberate provisional-sample mismatch exposed the wrong saved score, and a hidden-row comparison exposed an unwanted top-N entry.
fork/59	A coverage audit identified the remaining field-policy and supported memory-residency interactions not owned by focused bug regressions.

Dependency and scope

This branch has merge base 2371e9022a3d2a046d9b70efa17efa8874b7e9a7 with upstream main. The 11 production-fix deltas from fork/47–49 and fork/51–58 were replayed as signed commits in dependency order, and fork/59 contributes test-only coverage. Current-main reconciliation commits only adapt tests and add no production behavior.

Split PR	Exact source head
fork/47	a1541f0b283798b6ef80c69684260fae202909f3
fork/48	2e3d163c29d2e9cadde5f23243521e4986999bc8
fork/49	459bcc30f01a1e762bf31cf5e435a78662533794
fork/51	f666c70e905f411300b929d7913e3567c7013228
fork/52	b904d94f01eb4afa8c0fa4aa11628af78c857f99
fork/53	4a68479db6d5c20c4ebdf8f60fcf7ca664c004b3
fork/54	10238201b2f30b3782cb8f8e2fcf1a5cce26a99c
fork/55	7df61bd86d74847f769eeed7429583463f4ca58f
fork/56	7024cdb8adb8000cc9da19343b9f18b22ba54ad5
fork/57	8e21607ced43a2ee46827d7e1f17743d8d359b60
fork/58	f38fce45358a561e748cf4174b4af7feef20709c
fork/59	29bb8f9f7e81106ed4c15924dca0889e54b5b0ff

The fork PRs remain the detailed review records; this PR is their upstream landing path, so the same changes should not later be landed from both places. Fork PR #50 is intentionally excluded because NVIDIA/Megatron-LM #7063 owns and supersedes that behavior.

The aggregate also excludes the closed fork/40 CUDA-graph prerequisite and #59's unsupported UVM+OFFLOAD row. The two supported #59 runtime rows remain, and its exhaustive policy manifest includes fields supplied by prerequisite work already merged through #7031 and #7053. These current-main reconciliations are test-only.

The change adds no functional-test files. Its unit-test delta is 2,339 additions and 340 removals, or 1,999 net lines, within the campaign's 2,000-line budget.

Validation

• git diff --check passed against current upstream main.
• Black, isort, Pylint, Ruff, and py_compile passed for the changed Python surface.
• The field-policy audit covers all 40 request fields and all 19 sampling fields with no missing or stale entries.
• Exact aggregate head: 9df76ecadb6f5108e7bfdce47e93b04c60422db0.
• At the exact head, targeted admission, lifecycle, and downstream-consumer validation passed across eight ranks: 328 passes, 16 expected skips, and zero failures or deselections.
• The full lifecycle matrix passed on every rank at the production-identical predecessor: 125 passes and 5 expected skips per rank. Its separate two-rank expert-parallel case passed on both ranks.
• A broader API and caller sweep also passed every non-baseline test at a production-identical predecessor: 391 passes and 12 expected skips per rank.
• Two GDN tests in the broader engine sweep fail identically on this branch and unmodified current main; they are a pre-existing main issue outside this change.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tdene]

I closely reviewed each of the individual PRs, and they were all good. Approved.

[tdene]

/ok to test 9df76ec