[lmcafee-nvidia]

What does this PR do?

October 20 main refresh: head 3d8fb014ff90f3dd72751a7a054bebcf97f03b9d, parent 184ed356a80030f38f448e60ff0ebe1f622f9d5d, pinned NVIDIA main 474b1e9ee72fbe19205a6dfe26d56620977a6275. Existing branches and dependencies are preserved. Shared fixtures select their intended attention backend explicitly and observe the current async forward state.

This PR adds the shared unit-test foundation for data-parallel inference coordination. Real clients submit requests through the coordinator to real model replicas; the tests check that routing preserves model results while the selected inference behavior actually executes.

The foundation adds:

• one shared coordinator/client/model harness with direct-engine references, request-owned runtime observations, and normal teardown;
• 97 authored selections across routing, sampling, model topology and kernels, lifecycle, media, scheduling, batch invariance, backends, execution controls, and admission; and
• deliberate workloads for cache reuse, overlapping prefill/decode, metadata compaction, active suspension, allocator pressure, and different physical batch sizes.

The 97 selections include parameter variants and correctness controls. They are not 97 distinct positive pairwise cells. Production fixes and their focused regressions have separate owners in the split.

How was the test matrix determined?

The source audit found a gap between coordinator plumbing coverage and real inference behavior. Existing coordinator tests use a DummyEngine with synthetic progress, or fake rank identities and seeded hashes to check routing choices. Those tests cannot establish that the chosen replica consumes real cached KV, executes a model kernel, preserves a sampling distribution, or resumes a partly generated request correctly.

The matrix follows the production boundaries exposed by coordinator routing and handlers, engine configuration, request sampling fields, and model/parallel execution. It has three layers:

1. Routing and request-result contracts: overlapping client-local IDs, streaming, sampling policies, output opt-ins, BOS/stop handling, cache-owner selection, and rejection retirement.
2. Real execution interactions: model-parallel owners, graph replay, MTP, recurrent state, attention/precision backends, media embedding consumption, and batch-shape changes.
3. Stateful schedules: partial prefill beside decode, affinity versus load/expiry, suspension and residency, cancellation, epochs, allocator pressure, and idle expert progress.

Deterministic rows compare direct and routed execution using the same engine weights and companion configuration. The stochastic rows compare the first-step logits and filtered distribution at the same prompt/history, then check the sampled token's reported score; they do not require identical random draws. Stateful rows additionally assert the exact transition or ownership invariant. Each claimed interaction needs runtime evidence from the affected routed work.

This is a bounded differential interaction matrix. It does not enumerate every CLI combination, every companion-disabled cell, or every fault phase. Parameter products are expanded only where listed below; backend wrappers reuse the all-owner oracle, and graph-selection rows reuse the chunk-partition workload without adding another copy of either helper. Static enumeration of the frozen source agrees with the extraction author's inventory; it is not a runtime collection result.

What combinations are tested?

All paths below are under tests/unit_tests/inference/. Counts are unique authored selections, counted once regardless of rank count or repeated consumer runs:

8 + 28 + 13 + 12 + 12 + 4 + 1 + 6 + 12 + 1 = 97.

Routing, streaming, and prefix ownership (8 selections)

test_coordinator_pairwise.py uses real clients with overlapping local request ID zero. Ordinary output rows submit one request per DP owner; prefix rows warm an owner and repeat a 33-token prompt with 16-token KV blocks and LRU eviction.

Test / configuration	Actual cases	Count
test_routed_output_matches_direct, final result	LEGACY / interval 0; ASYNC / interval 0	2
Same test, streaming every token	LEGACY / interval 1; ASYNC / interval 1	2
Same test, streaming every three tokens	LEGACY / interval 3; ASYNC / interval 3	2
test_routed_prefix_reuse_executes_on_warm_owner	LONGEST_PREFIX; FIRST_PREFIX_BLOCK	2

Sampling and request output policy (28 selections)

test_coordinator_sampling_pairwise.py contributes 20 stochastic selections, four output opt-ins, and four BOS/stop-compaction selections. Stochastic requests generate one token at temperature 0.7; each backend/filter row runs all four combinations of shared versus DP-offset seed and raw versus processed log probabilities.

Backend/filter case	(top_k, top_p)	Seed × score modes for each row	Count
torch-top-k	(4, 0.0)	shared/raw; shared/processed; rank/raw; rank/processed	4
torch-top-p	(0, 0.7)	shared/raw; shared/processed; rank/raw; rank/processed	4
flashinfer-top-k	(4, 0.0)	shared/raw; shared/processed; rank/raw; rank/processed	4
flashinfer-top-p	(0, 0.7)	shared/raw; shared/processed; rank/raw; rank/processed	4
flashinfer-joint	(4, 0.7)	shared/raw; shared/processed; rank/raw; rank/processed	4

Here shared/rank means offset_sampling_seed_by_dp_rank=False/True, and raw/processed means raw_logprobs/processed_logprobs.

test_routed_output_opt_ins: (echo, detokenize, eos)	Expected result	Count
(False, False, False)	Five generated tokens; prompt/text omitted	1
(True, True, False)	Five generated tokens; prompt and decoded text returned	1
(True, False, False)	Five generated tokens; prompt returned, text omitted	1
(False, False, True)	Terminate on the direct reference's first token	1

test_routed_bos_stop_after_metadata_compaction: (keep, skip)	Stop tokens / prompt scores	Count
(False, False)	Remove matching stop; retain prompt scores	1
(False, True)	Remove matching stop; skip prompt scores	1
(True, False)	Keep matching stop; retain prompt scores	1
(True, True)	Keep matching stop; skip prompt scores	1

The last four cases insert BOS, place a stochastic short request before the greedy target on its owner, and require the target's top_k=1 metadata to survive compaction after the short request retires.

Model topology and feature execution (13 selections)

test_coordinator_features_pairwise.py compares direct output with routed output on every DP owner and every model-parallel role. The three topology rows use eight ranks; ordinary feature rows use two DP replicas. Each row below is one selection.

Actual case	Configuration exercised	Required execution evidence
tp2-dp4	TP2, PP1, DP4, local modules	Real TP partitions and reductions; every owner/role prefills and decodes
pp2-dp4	TP1, PP2, DP4, local modules	Both pipeline roles execute each owner's request through prefill and decode
tp2-pp2-sp-dp2	TP2 × PP2 × DP2, sequence parallelism	Real SP gather/reduce-scatter and their tensor dimensions
test_routed_cuda_graph_replays_for_every_target	Four graphs, decode replay, maximum 16 graph tokens	Completed production replay on active routed targets
test_routed_mtp_executes_inner_steps_on_every_target	One speculative token, six generated tokens	Real inner MTP step, speculative steps, and proposals
Hybrid mamba	Hybrid model with Mamba mixer	Target-mapped convolution and recurrent state change during decode
Hybrid gdp	Hybrid model with Gated Delta Product mixer	Same owned-state invariant and actual GDP mixer execution
Hybrid gdn	Hybrid model with Gated Delta Net mixer	Same owned-state invariant and actual GDN mixer execution
fp8	FP8, hidden size 128	Quantized forwards and finite sampled logits
fused-rope	Native FlashInfer NeoX D64, hidden size 256	Real fused rotary call and finite sampled logits
swa-sink	Window (4, 0), off-by-one softmax	Sliding-window kernel and sink correction
alternating-swa-sink	Window (4, 0), skip frequency 2, learnable softmax	Sliding-window and full-attention kernels plus sink correction
test_routed_moe_participates_in_ep_collectives	EP2, real MoE, NCCL dispatcher, inference-optimized modules	Balanced dispatch/combine, no ordering violation or outstanding dispatch

Request lifecycle and memory residency (12 selections)

test_coordinator_lifecycle_pairwise.py pauses an active request only after real generation has begun. Residency rows compare the resumed 16-token result with uninterrupted direct execution; the other schedules exercise cancellation, epoch propagation, pressure, and drained reset.

Actual cases	Deliberate transition	Count
LEGACY/PERSIST; ASYNC/PERSIST	Suspend/resume active work while preserving the KV allocation pointer	2
LEGACY/OFFLOAD; ASYNC/OFFLOAD	Native memory-saver region/pause/resume, exact target KV restoration, stable pointer, subsequent graph replay	2
LEGACY/RECOMPUTE; ASYNC/RECOMPUTE	Drop KV allocation, invalidate pending async logits, and prefill the target again	2
Abort queued; abort active	Cancel client A's local ID zero while client B's local ID zero survives	2
Epoch update with synchronous ZMQ collectives False; True	Reach epoch 7 on active and idle engines and preserve returned epoch records	2
Allocator-pressure victim identity	Five-block pool, one-block paused limit, three requests per owner; pause/evict the same ID and observe its repeated prefill	1
UVM drained reset	Real managed allocation extent and CPU/GPU prefetch; retain live control objects and serve another request after reset	1

Media routing, embedding reuse, and epochs (12 selections)

test_coordinator_media_pairwise.py uses two replicas and real LLaVA vision, projection, and language execution. Each row runs both allow_stale_multimodal_embeddings=False (invalidate) and True (retain) across suspend/resume. Prefix routing is LOAD_BALANCED, prefix eviction is LRU, and the vision cache is enabled.

Media case	Media coordinator policy	Actual stale-cache cases	Count
clip, fixed image	AFFINITY	invalidate; retain	2
clip, fixed image	LOAD_BALANCED	invalidate; retain	2
radio-dynamic, dynamic image	AFFINITY	invalidate; retain	2
radio-dynamic, dynamic image	LOAD_BALANCED	invalidate; retain	2
radio-video, two frames	AFFINITY	invalidate; retain	2
radio-video, two frames	LOAD_BALANCED	invalidate; retain	2

A competing text request puts the warm image on its peer. After the epoch transition, the workload repeats the same media and then submits different media, checking actual owner choice, vision/projection call counts, embedding injection at 4/6/12 image positions, cache salt changes, and zero incorrect prefix reuse for different media. It checks the configured cache policy across an epoch transition without changing model weights.

Chunk scheduling and routing priorities (4 selections)

test_coordinator_schedule_pairwise.py uses two replicas and records work on the selected target owner.

Actual case	Schedule/configuration	Count
test_routed_chunk_partition_beside_decode	65-token target, 32-token budget, at least three exact prompt partitions, one sharing a step with decode	1
affinity	Warm 33-token prompt; routing alpha 0, TTL 300; warm owner wins despite competing load and actually reuses KV	1
load	Same competing schedule; alpha 100, TTL 300; another owner wins and target has no warm KV reuse	1
expired	Alpha 0, TTL 0; expired routing affinity sends the target away from the warm owner	1

Batch invariance at a physical padding boundary (1 selection)

test_coordinator_batch_invariant_pairwise.py uses native Transformer Engine batch invariance, ordered collectives, FA3, BF16, and legacy scheduling. The shared factory supplies the real model.

Actual case	Direct → routed comparison	Count
test_routed_batch_invariance_at_padding_boundary	One-token prompt alone → 65 requests per owner, 130 replies across DP2; real GEMM physical rows grow from 64 to 128	1

The selected request's input tokens, positions 0, 1, 2, 3, and finite logits must be bitwise identical across the two schedules; every routed result also matches the direct output.

Dense backends and native model modes (6 selections)

test_coordinator_backends_pairwise.py reuses the all-owner kernel test, including direct/routed result parity and finite sampled logits. Each row below is one selection.

Actual case	Configuration	Required native witness
optimized	Inference-optimized modules, FA3 configuration	Real inference-optimized module forward
fa2	Local modules, 256-token blocks/budget	FA2 prefill and decode kernel calls
fa4	Local modules, hidden size 128	FA4 prefill and decode kernel calls
mla	Latent KV cache, 64 heads, QK normalization, RoPE, FA3 configuration	FlashMLA decode with latent width 576 and value width 512
mxfp8	Native TE MXFP8, FA4, hidden size 256	Actual MXFP8 GEMM operands in prefill and decode, physical batch 128
nvfp4	Native TE NVFP4, FA4, hidden size 256	Actual NVFP4 GEMM operands in prefill and decode, physical batch 128

Native precision uses the same 128-row physical batch for direct and routed runs. These two rows check ordinary precision/routing parity, independently of the batch-invariance row's deliberate shape change.

Graph selection and idle expert progress (12 selections)

test_coordinator_execution_pairwise.py adds six graph-sizing selections and six idle-progress selections. Graph rows reuse the 65-token chunked target with three graphs, mixed-prefill count 2, and maximum decode graph size 16; they check actual replay dimensions and retain the chunk-partition oracle.

Graph sizing distribution	Actual cuda_graph_all_prefills cases	Expected graph tokens for the observed seven-token mixed step	Count
EXPONENTIAL	False; True	8; 16	2
LINEAR	False; True	12; 12	2
HYBRID	False; True	8; 16	2

All observed dimensions must belong to the captured set and match physical input size. A target-bearing pure-prefill replay above 16 tokens must occur exactly when all-prefill capture is enabled.

Idle-progress configuration	Actual synchronous-ZMQ cases	Companion behavior	Count
Consensus interval 20, EP2	False; True	Idle peer participates in NCCL work covering the busy owner's dispatches; observed cadence is divisible by 20	2
Consensus interval 3, EP2	False; True	Same ownership/dispatch invariant with cadence divisible by 3	2
Interval parameter None, consensus disabled, EP1	False; True	Independent replicas still make progress; dummy forwards occur and no consensus call occurs	2

These rows generate 32 tokens on one owner, compare token/status output with the direct reference, require finite score arrays, and check balanced NCCL dispatch/combine where EP2 applies. The disabled-consensus EP1 cases do not claim disabled-consensus EP2 support.

Rejected admission and subsequent progress (1 selection)

test_coordinator_admission_pairwise.py checks rejection before real model execution, then sends a healthy request to the same owner.

Actual case	Request sequence and oracle	Count
test_rejected_admission_retires_before_healthy	A 257-token prompt exceeds the 256-token limit: one FAILED terminal with MaxSequenceLengthOverflowError, no forward or retained request state; the subsequent healthy request matches direct output	1

This selection is a rejection/liveness control and receives no positive model-feature pair credit.

How are results checked?

The harness observes completed production forwards and coordinator handler events. Tests force assignment while engines are paused, record the actual owner and client-local IDs, then release real execution. The all-owner helper gathers records across ranks and requires every DP owner and model-parallel role to prefill and decode its assigned IDs, with disjoint ownership between replicas. Warmup, direct-reference calls, and bookkeeping-only work cannot satisfy that helper's target evidence.

The oracles match each contract:

• Deterministic all-owner rows require exact status, generated tokens, prompt tokens, and generated text against the direct reference. Basic routing additionally compares prompt/generated log probabilities with torch.testing.assert_close defaults, and requires incremental partials, aligned token/selected-score lengths, exactly one final stream item, and no duplicate terminal engine reply.
• Sampling requires exact first-step logits and filtered distributions (rtol=0, atol=0), RNG consumption, the expected per-owner seed, and a real FlashInfer joint-filter call where selected. The returned selected-token log probability uses absolute tolerance 8 * finfo(dtype).eps and zero relative tolerance; raw and processed scores must differ and the sampled token must remain in the nontrivial filtered support.
• BOS/stop rows compare prompt scores, require matching generated-token/selected-score/top-N lengths, and check selected scores against their top-N entries after observed metadata compaction. The separate streamed top-N holdback regression is excluded from this foundation.
• Cache, lifecycle, and media rows require actual cache consumption, target-owned transition evidence, or exact memory/state invariants. OFFLOAD checks native memory-saver calls and bitwise target KV restoration; media checks the projected embeddings actually injected into the language model; chunking accounts for every prompt token and requires overlap with decode.
• Kernel rows require finite logits and the named target-attributed native calls. Batch invariance additionally requires bitwise target-logit equality across real 64→128 GEMM batches. Idle-progress rows require finite generated-score arrays; they do not assert numerical log-probability parity between different batch schedules.

Retirement is checked before teardown: engine requests/queues and active context rows are empty, coordinator request maps and pending counts are empty, client futures/submission timestamps/streams are empty, and the coordinator remains alive after ordinary completion. The shared fixture then performs the original pause/stop/barrier/shutdown sequence, joins the coordinator, closes clients and sockets, deletes graphs, and destroys model-parallel state. Original assertions, parameter cases, timeouts, fixtures, and teardown are retained.

Scope

This PR adds 11 unit-test/helper files, +2185/-0 lines, and has no production delta or functional-test/config/runner changes. It owns 97 selections; the 18 focused child selections remain in #82/#86/#87/#83/#85. Together they preserve the 115-selection campaign.

The current candidate is 79e2160ee2d5f8458dd6184e9de46470edc40c7b above dp-prerequisites-v2 at 5a1c4da4d27b297078f8fe74572731d556ea39a6. Main 4729b9c6cc17af2a1b47bf7c42b05059f36ca3b9 was merged into the campaign-owned prerequisite branch first, then propagated through the foundation and descendants with signed append-only merges.

The campaign is self-contained: its own prerequisite branch copies the FA3 trailing-return-metadata fix and the TE norm/projection source-device fix; it does not merge peer campaign branches. Four earlier copied prerequisite concepts are now supplied by upstream NVIDIA#7256 and no longer appear as production differences from main.

The test source originated in historical fork/76, which remains frozen. Current API adaptations consume flat completed-request futures, decode prompt-token tensors through production deserialization, and check coordinator-detokenized text against the direct generated-token oracle, and inspect the current vision-cache entry.embedding payload. All original activation witnesses, parity/invariant checks, parameters, timeouts and normal teardown remain.

Remaining copied prerequisite	Original replay	Foundation consumer
FA3 trailing return metadata	0865a7b01ab0e3a062f60ccc305afbf48fe45faa	Sink attention
TE split-module source device	55fb75d133dfdde2f6b3e347c4ce09063c69c451	Native MLA

New main features, including multi-EOS defaults, payload staging, cached MTP drafts, nondefault softcap, BI+MXFP8 MoE and dynamic CP, are not automatically covered by existing separate rows. Runtime configuration/plumbing alone receives no pairwise credit. These stable fork review records remain the primary records; later upstream landing uses new linked PRs.

Validation

October 20 refresh acceptance

Candidate 3d8fb014ff90f3dd72751a7a054bebcf97f03b9d includes pinned main 474b1e9ee72fbe19205a6dfe26d56620977a6275.

DFW jobs 20558189, 20563746: all 10 required groups passed on every required rank, including setup, execution and teardown.

HSG jobs 8095236, 8095800: all 2 required groups passed on every required rank, including setup, execution and teardown.

No selected case was skipped.

Validation uses qualified compatibility environments: DFW H100 with Torch 2.10/CUDA 13.1, TE 2.20 and FlashInfer 0.6.16.post3; where listed, HSG GB200 with Torch 2.11/CUDA 13.2, TE 2.20 and FlashInfer 0.6.8. This does not claim the exact current 26.09 CI image.

Signed append-only ancestry, Python compilation, whitespace, Black, isort, Ruff, Pylint and the actual-parent kernel gate passed. Existing documented non-numeric kernel-gate exemptions remain on #65 and #102. No unit tests ran on the local workstation.

Earlier validation records

Tested candidate: 79e2160ee2d5f8458dd6184e9de46470edc40c7b (tree 3899970c34ba7eafaf4cd10acb9b83c6115f3b2e).

All 97 unique authored foundation selections passed: DFW jobs 19788479 and 19788480 cover ordinary, media, lifecycle and native backend rows; HSG 7686748 covers MXFP8 and NVFP4. Every two-rank row passed on both ranks and each of the three topology rows passed on all eight ranks. All authored selections have zero skips.

The complete refreshed union passes all 115 unique authored selections, with zero authored skips, and remains at the approved 2,630 net unit-test-line cap. Current-main production oracles pass five stop/admission cases in DFW 19787958 and two pressure/reset cases in 19788480; source checks confirm zero production differences from pinned main.

The DFW compatibility environment uses TE 2.20.0+131157b5, FlashInfer 0.6.16.post3 and Mamba 2.3.2.post1 on Torch 2.10/CUDA 13.1. Native Blackwell precision uses TE 2.20 on HSG with Torch 2.11/CUDA 13.2. These qualified native environments are not a byte-for-byte reproduction of the 26.08 lock environment.

Pinned isort 5.13.2, Black 26.3.0, Ruff 0.12.11, AST, diff checks, signed/DCO history, append-only ancestry, parent-relative scope and the kernel-determinism gate passed. Post-run source receipts confirm clean exact commits. No unit tests, collection or test imports ran locally.

[lmcafee-nvidia]

hi @sidsingh-nvidia, can you please review this DP coordinator test foundation? thanks!