From 396b7720f699f7137db6b085dd1284b8ac209208 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 14 Sep 2026 11:48:26 -0400 Subject: [PATCH 1/9] Add DP coordinator pairwise test foundation Signed-off-by: Lawrence McAfee --- .../inference/coordinator_pairwise_utils.py | 364 +++++++++++++++++ .../test_coordinator_admission_pairwise.py | 48 +++ .../test_coordinator_backends_pairwise.py | 133 +++++++ ...st_coordinator_batch_invariant_pairwise.py | 82 ++++ .../test_coordinator_execution_pairwise.py | 140 +++++++ .../test_coordinator_features_pairwise.py | 372 ++++++++++++++++++ .../test_coordinator_lifecycle_pairwise.py | 284 +++++++++++++ .../test_coordinator_media_pairwise.py | 309 +++++++++++++++ .../inference/test_coordinator_pairwise.py | 126 ++++++ .../test_coordinator_sampling_pairwise.py | 205 ++++++++++ .../test_coordinator_schedule_pairwise.py | 120 ++++++ 11 files changed, 2183 insertions(+) create mode 100644 tests/unit_tests/inference/coordinator_pairwise_utils.py create mode 100644 tests/unit_tests/inference/test_coordinator_admission_pairwise.py create mode 100644 tests/unit_tests/inference/test_coordinator_backends_pairwise.py create mode 100644 tests/unit_tests/inference/test_coordinator_batch_invariant_pairwise.py create mode 100644 tests/unit_tests/inference/test_coordinator_execution_pairwise.py create mode 100644 tests/unit_tests/inference/test_coordinator_features_pairwise.py create mode 100644 tests/unit_tests/inference/test_coordinator_lifecycle_pairwise.py create mode 100644 tests/unit_tests/inference/test_coordinator_media_pairwise.py create mode 100644 tests/unit_tests/inference/test_coordinator_pairwise.py create mode 100644 tests/unit_tests/inference/test_coordinator_sampling_pairwise.py create mode 100644 tests/unit_tests/inference/test_coordinator_schedule_pairwise.py diff --git a/tests/unit_tests/inference/coordinator_pairwise_utils.py b/tests/unit_tests/inference/coordinator_pairwise_utils.py new file mode 100644 index 00000000000..598c0d90733 --- /dev/null +++ b/tests/unit_tests/inference/coordinator_pairwise_utils.py @@ -0,0 +1,364 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Real-model DP routing fixtures; observations never replace production behavior.""" + +import asyncio +import copy +import queue +import threading +from contextlib import asynccontextmanager +from types import SimpleNamespace + +import msgpack +import torch +import zmq + +from megatron.core.inference.data_parallel_inference_coordinator import ( + DataParallelInferenceCoordinator, +) +from megatron.core.inference.engines.async_zmq_communicator import AsyncZMQCommunicator +from megatron.core.inference.engines.dynamic_engine import EngineState +from megatron.core.inference.headers import Headers +from megatron.core.inference.inference_client import InferenceClient +from megatron.core.inference.sampling_params import SamplingParams +from megatron.core.transformer.cuda_graphs import delete_cuda_graphs +from megatron.core.transformer.enums import AttnBackend +from megatron.core.transformer.module import Float16Module +from tests.unit_tests.inference.engines import test_dynamic_engine as engine_tests +from tests.unit_tests.inference.test_data_parallel_inference_coordinator import DummyTokenizer +from tests.unit_tests.test_utilities import Utils + + +async def until(predicate, timeout=30): + """Wait for an observed state, not for a presumed scheduling delay.""" + + async def poll(): + while not predicate(): + await asyncio.sleep(0) + + await asyncio.wait_for(poll(), timeout) + + +class CoordinatorThread: + """Run the unmodified blocking coordinator in its own socket-owning thread.""" + + def __init__(self, dp_size, tokenizer, *, setup=None, **options): + self.events = [] + self.errors = [] + self.addresses = queue.Queue() + self.ready = threading.Event() + self.coordinator = None + self.thread = threading.Thread( + target=self._run, args=(dp_size, tokenizer, setup, options), daemon=True + ) + self.thread.start() + self.address = self.addresses.get(timeout=30) + + def _run(self, dp_size, tokenizer, setup, options): + try: + coordinator = DataParallelInferenceCoordinator( + pipe_connection=SimpleNamespace(send=self.addresses.put, close=lambda: None), + data_parallel_size=dp_size, + tokenizer=tokenizer, + max_requests=8, + deterministic_mode=True, + hostname="127.0.0.1", + **options, + ) + self.coordinator = coordinator + coordinator.router_socket.setsockopt(zmq.LINGER, 0) + if setup is not None: + setup(coordinator) + for header, handler in list(coordinator._handlers.items()): + + def observe(c, sender, metadata, bodies, handler=handler, header=header): + before = dict(c.request_id_to_rank) + result = handler(c, sender, metadata, bodies) + self.events.append( + dict( + header=header, + sender=sender, + metadata=metadata, + before=before, + after=dict(c.request_id_to_rank), + pending=c._pending_counts.tolist(), + ) + ) + return result + + coordinator._handlers[header] = observe + self.ready.set() + coordinator.start() + except BaseException as error: + self.errors.append(error) + finally: + if self.coordinator is not None: + self.coordinator.router_socket.disable_monitor() + self.coordinator.stop() + + def assert_requests_retired(self): + c = self.coordinator + assert not self.errors, self.errors + for name in ( + "request_id_to_client_id", + "request_id_to_client_request_id", + "client_request_to_request_id", + "request_id_to_rank", + ): + assert not getattr(c, name), name + assert not getattr(c, "request_id_to_sampling_params", {}) + assert not c._pending_counts.any() + assert self.thread.is_alive(), "Ordinary completion must not stop the service" + + +class RoutedModel: + """One real engine per model-parallel group plus rank-zero client controls.""" + + def __init__(self, config, engine_factory=None, **coordinator_options): + self.config = config + self.rank = torch.distributed.get_rank() + self.dp_size = Utils.world_size // ( + config.tensor_model_parallel_size * config.pipeline_model_parallel_size + ) + assert self.dp_size >= 2, "Routing rows require at least two DP replicas" + self.engine = ( + engine_factory(config) + if engine_factory is not None + else engine_tests.DynamicInferenceEngineTestBase._build_test_env(config).engine + ) + self.tokenizer = DummyTokenizer(vocab_size=config.vocab_size, eod=-1) + self.engine.controller.tokenizer = self.tokenizer + self.coordinator_options = coordinator_options + self.clients = [] + self.witnesses = [] + self.service = None + self.sync_context = zmq.Context() + self.sync = AsyncZMQCommunicator( + self.sync_context, process_group=None, hostname="127.0.0.1" + ) + + async def barrier(self): + await asyncio.wait_for(self.sync.all_reduce_max(1), timeout=60) + + async def direct(self, prompt, params): + """Same weights, topology and one-request schedule, without routing.""" + future = self.engine.add_request(10001, prompt, copy.deepcopy(params)) + while self.engine.has_unfinished_requests(): + await self.engine.async_step() + result = msgpack.unpackb( + msgpack.packb((await future).merge().serialize(), use_bin_type=True), raw=False + ) + self.engine.reset() + return result + + async def start(self): + if self.rank == 0: + self.service = CoordinatorThread( + self.dp_size, + self.tokenizer, + block_size_tokens=self.config.context_block_size_tokens, + enable_prefix_caching=self.config.enable_prefix_caching, + **self.coordinator_options, + ) + address = [self.service.address if self.service else None] + torch.distributed.broadcast_object_list(address, src=0) + self.address = address[0] + self._observe_forwards() + await self.engine.start_listening_to_data_parallel_coordinator( + inference_coordinator_port=int(self.address.rsplit(":", 1)[-1]), + launch_inference_coordinator=False, + hostname="127.0.0.1", + ) + if self.rank == 0: + await until(self.service.ready.is_set) + assert len(self.service.coordinator.identities_of_data_parallel_ranks) == self.dp_size + for _ in range(self.dp_size): + client = InferenceClient(self.address) + client.start(connect_timeout_seconds=30) + self.clients.append(client) + await self.barrier() + + def _observe_forwards(self): + model = self.engine.controller.inference_wrapped_model.model + forward = model.forward + + def observed_forward(*args, **kwargs): + context = self.engine.context + ids = [ + request_id + for request_id in context.request_ids[ + context.paused_request_count : context.total_request_count + ].tolist() + if request_id in self.engine.requests + ] + snapshot = dict( + ids=ids, + prefill=context.num_prefill_requests, + decode=context.num_decode_requests, + chunk=context.chunked_prefill_request_id, + tokens=context.active_token_count, + padded_tokens=context.padded_active_token_count, + async_step=context.async_sched_step_count, + cached={rid: self.engine.get_request(rid).num_cached_tokens for rid in ids}, + ) + result = forward(*args, **kwargs) + # Commit evidence only after the actual model forward returns. + if ids: + self.witnesses.append(snapshot) + return result + + model.forward = observed_forward + + async def pause(self): + if self.rank == 0: + self.clients[0].pause_engines() + await asyncio.wait_for(self.engine.wait_until(EngineState.PAUSED), timeout=60) + await self.barrier() + + async def unpause(self): + if self.rank == 0: + self.clients[0].unpause_engines() + await asyncio.wait_for(self.engine.wait_until(EngineState.RUNNING), timeout=60) + await self.barrier() + + async def run_paused_batch(self, submissions): + """Return replies, actual owners and client-local IDs for an ordered batch.""" + await self.pause() + payload = [None] + if self.rank == 0: + pending, local_ids = [], [] + coordinator = self.service.coordinator + for client, prompt, params, *media in submissions: + request_id, future = self.clients[client].add_request_with_id( + prompt, copy.deepcopy(params), multi_modal_data=media[0] if media else None + ) + pending.append(future) + local_ids.append((client, request_id)) + await until(lambda: len(coordinator.request_id_to_rank) == len(pending)) + owners = dict(coordinator.request_id_to_rank) + assert all( + coordinator._pending_counts[index] == list(owners.values()).count(identity) + for identity, index in coordinator.identity_to_rank_index.items() + ) + await self.unpause() + if self.rank == 0: + payload[0] = (await asyncio.wait_for(asyncio.gather(*pending), 120), owners, local_ids) + await self.barrier() + torch.distributed.broadcast_object_list(payload, src=0) + return payload[0] + + def assert_retired(self): + assert not self.engine.requests + assert not self.engine.waiting_request_ids + assert self.engine.context.get_active_request_count() == 0 + if self.rank == 0: + self.service.assert_requests_retired() + for client in self.clients: + assert not client.completion_futures + assert not client.request_submission_times + assert not client.streams + + async def close(self): + if hasattr(self.engine, "engine_loop_task"): + await self.pause() + if self.rank == 0: + self.clients[0].stop_engines() + await asyncio.wait_for(self.engine.engine_loop_task, timeout=60) + await self.barrier() + if self.rank == 0 and self.service is not None: + if self.clients: + self.clients[0].shutdown_coordinator() + await asyncio.to_thread(self.service.thread.join, 30) + assert not self.service.thread.is_alive() + assert not self.service.errors, self.service.errors + for client in self.clients: + client.stop() + await asyncio.gather(*(c.listener_task for c in self.clients), return_exceptions=True) + self.sync.close() + self.sync_context.term() + + +@asynccontextmanager +async def routed_model(monkeypatch, *, coordinator_options=None, engine_factory=None, **overrides): + """Build a real distributed engine, optionally via a config-to-engine factory. + + Backend globals must be selected by the caller before entering this fixture. + """ + monkeypatch.setenv("CUDA_DEVICE_MAX_CONNECTIONS", "1") + settings = dict( + num_requests=0, + max_sequence_length=256, + context_max_requests=8, + context_max_tokens=128, + context_block_size_tokens=16, + context_buffer_size_gb=0.02, + transformer_impl="transformer_engine", + ) + settings.update(overrides) + config = engine_tests.DynamicEngineTestConfig(**settings) + Utils.initialize_model_parallel( + tensor_model_parallel_size=config.tensor_model_parallel_size, + pipeline_model_parallel_size=config.pipeline_model_parallel_size, + expert_model_parallel_size=config.expert_model_parallel_size, + ) + harness = None + try: + with torch.inference_mode(): + harness = RoutedModel(config, engine_factory, **(coordinator_options or {})) + yield harness + finally: + if harness is not None: + await harness.close() + delete_cuda_graphs() + Utils.destroy_model_parallel() + + +def greedy_params(**overrides): + params = dict(num_tokens_to_generate=8, termination_id=-1, top_k=1) + params.update(overrides) + return SamplingParams(**params) + + +def batch_invariant_engine(config): + """Construct native-TE BI modules; caller pins backend and rounders before CUDA.""" + assert config.tensor_model_parallel_size == config.pipeline_model_parallel_size == 1 + torch.manual_seed(config.random_seed) + engine_tests.model_parallel_cuda_manual_seed( + config.random_seed, inference_rng_tracker=True, force_reset_rng=True + ) + model_config = engine_tests.TransformerConfig( + num_layers=2, + hidden_size=128, + num_attention_heads=4, + use_cpu_initialization=True, + hidden_dropout=0.0, + attention_dropout=0.0, + batch_invariant_mode=True, + batch_invariant_backend="te_native", + flash_attention_version=3, + normalization="RMSNorm", + params_dtype=torch.bfloat16, + bf16=True, + attention_backend=AttnBackend.flash, + transformer_impl="transformer_engine", + nccl_all_reduce_for_prefill=False, + inference_rng_tracker=True, + inference_sampling_seed=config.random_seed, + ) + model = ( + engine_tests.GPTModel( + config=model_config, + transformer_layer_spec=engine_tests.get_gpt_layer_with_transformer_engine_spec(), + vocab_size=config.vocab_size, + max_sequence_length=config.max_sequence_length, + position_embedding_type="rope", + ) + .cuda() + .eval() + ) + context = engine_tests.DynamicInferenceEngineTestBase._build_inference_context( + config, model_config, [] + ) + wrapper = engine_tests.GPTInferenceWrapper(Float16Module(model_config, model).eval(), context) + controller = engine_tests.TextGenerationController(wrapper, DummyTokenizer(config.vocab_size)) + return engine_tests.DynamicInferenceEngine(controller, context) diff --git a/tests/unit_tests/inference/test_coordinator_admission_pairwise.py b/tests/unit_tests/inference/test_coordinator_admission_pairwise.py new file mode 100644 index 00000000000..33fec58e551 --- /dev/null +++ b/tests/unit_tests/inference/test_coordinator_admission_pairwise.py @@ -0,0 +1,48 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +from asyncio import wait_for + +import pytest + +from megatron.core.inference.headers import Headers +from tests.unit_tests.inference.coordinator_pairwise_utils import greedy_params, routed_model + +pytestmark = [pytest.mark.internal, pytest.mark.asyncio] + + +async def test_rejected_admission_retires_before_healthy(monkeypatch): + prompt, params = list(range(4, 20)), greedy_params() + async with routed_model(monkeypatch, max_sequence_length=256) as h: + direct = await h.direct(prompt, params) + await h.start() + bad = 0 + if h.rank == 0: + client, events = h.clients[0], h.service.events + failed = await wait_for(client.add_request([4] * 257, params), 60) + submitted = next(e["after"] for e in events if e["header"] == Headers.SUBMIT_REQUEST) + failed_id, owner = next(iter(submitted.items())) + bad = int((failed["request_id"], failed["status"]) != (failed_id, "FAILED")) + bad |= sum(e["type"] == "FAIL" for e in failed["events"]) != 1 + bad |= not any( + e["type"] == "ERROR_NONTRANSIENT" + and e["payload"]["type"] == "MaxSequenceLengthOverflowError" + for e in failed["events"] + ) + assert await h.sync.all_reduce_max( + len(h.engine.failed_request_ids), len(h.engine.requests), len(h.witnesses), bad + ) == (0, 0, 0, 0) + h.assert_retired() + if h.rank == 0: + final = await wait_for(client.add_request(prompt, params), 60) + for key in ("status", "generated_tokens", "generated_text"): + bad |= final[key] != direct[key] + submits = [e["after"] for e in events if e["header"] == Headers.SUBMIT_REQUEST] + bad |= submits[-1] != {final["request_id"]: owner} + replies = [ + item[0] + for e in events + if e["header"] == Headers.ENGINE_REPLY + for item in e["metadata"][1] + ] + bad |= replies.count(failed_id) != 1 + assert await h.sync.all_reduce_max(bad, int(bool(h.witnesses))) == (0, 1) + h.assert_retired() diff --git a/tests/unit_tests/inference/test_coordinator_backends_pairwise.py b/tests/unit_tests/inference/test_coordinator_backends_pairwise.py new file mode 100644 index 00000000000..a98a168b731 --- /dev/null +++ b/tests/unit_tests/inference/test_coordinator_backends_pairwise.py @@ -0,0 +1,133 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from functools import partial + +import pytest + +from megatron.core.transformer import attention +from megatron.core.transformer.transformer_config import MLATransformerConfig +from tests.unit_tests.inference import test_coordinator_features_pairwise as features +from tests.unit_tests.inference.engines import test_dynamic_engine as engines + + +@pytest.mark.internal +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("backend", "version", "symbols"), + [ + ("optimized", None, ()), + ("fa2", 2, ("flash_attn_varlen_func", "flash_attn_with_kvcache")), + ("fa4", 4, ("flash_attn4_varlen_func",)), + ], +) +async def test_routed_dense_backends_match_direct(monkeypatch, backend, version, symbols): + constructor = partial(engines.TransformerConfig, flash_attention_version=version or 3) + monkeypatch.setattr(engines, "TransformerConfig", constructor) + build = engines.DynamicInferenceEngineTestBase._build_test_env + routed = partial(features.routed_model, engine_factory=lambda config: build(config).engine) + monkeypatch.setattr(features, "routed_model", routed) + instrument = features._instrument_scenario_runtime + + def instrument_native(env, scenario, runtime): + instrument(env, scenario, runtime) + for symbol in symbols: + kernel = getattr(attention, symbol) + + def observed(*args, _kernel=kernel, **kwargs): + result = _kernel(*args, **kwargs) + phase = "decode" if env.engine.context.is_decode_only() else "prefill" + runtime[f"{backend}-{phase}"] += 1 + return result + + monkeypatch.setattr(attention, symbol, observed) + + monkeypatch.setattr(features, "_instrument_scenario_runtime", instrument_native) + options = {"transformer_impl": "inference_optimized" if version is None else "local"} + if version == 2: + options.update(context_block_size_tokens=256, context_max_tokens=256) + if version == 4: + options["hidden_size"] = 128 + signals = ("inference-optimized",) if version is None else () + required = ( + ("module-forward:inference-optimized",) + if version is None + else (f"{backend}-prefill", f"{backend}-decode") + ) + await features.test_routed_model_kernels_match_direct(monkeypatch, options, signals, required) + + +@pytest.mark.internal +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ["mla", "mxfp8", "nvfp4"]) +async def test_routed_native_model_modes_match_direct(monkeypatch, mode): + """One-request physical batches preserve ordinary, non-BI precision semantics.""" + constructor = MLATransformerConfig if mode == "mla" else engines.TransformerConfig + overrides = dict( + flash_attention_version=3 if mode == "mla" else 4, num_layers=2, add_bias_linear=False + ) + options = dict(hidden_size=256, context_max_requests=128, context_max_tokens=128) + if mode == "mla": + overrides.update( + multi_latent_attention=True, + cache_mla_latents=True, + num_attention_heads=64, + qk_layernorm=True, + ) + options.update(context_block_size_tokens=64, position_embedding_type="rope") + spec = engines.get_gpt_layer_with_transformer_engine_spec + monkeypatch.setattr( + engines, + "get_gpt_layer_with_transformer_engine_spec", + partial(spec, multi_latent_attention=True, qk_layernorm=True), + ) + else: + overrides.update( + dict(fp8="hybrid", fp8_recipe="mxfp8") + if mode == "mxfp8" + else dict(fp4="e2m1", fp4_recipe="nvfp4") + ) + # Both direct and routed GEMMs consume the same 128-row physical batch. + for name in ("ROUNDER", "TOKEN_ROUNDER", "REQUEST_ROUNDER"): + monkeypatch.setattr(engines.DynamicInferenceContext, name, 128) + monkeypatch.setattr(engines, "set_rounder", lambda _: None) + monkeypatch.setattr(engines, "TransformerConfig", lambda **kw: constructor(**(kw | overrides))) + build = engines.DynamicInferenceEngineTestBase._build_test_env + monkeypatch.setattr( + features, + "routed_model", + partial(features.routed_model, engine_factory=lambda c: build(c).engine), + ) + instrument = features._instrument_scenario_runtime + + def instrument_native(env, scenario, runtime): + instrument(env, scenario, runtime) + if mode == "mla": + sources = [(attention, "flash_mla_with_kvcache")] + else: + from transformer_engine.pytorch.module import layernorm_linear, linear + + sources = [(linear, "general_gemm"), (layernorm_linear, "general_gemm")] + for module, symbol in sources: + kernel = getattr(module, symbol) + + def observed(*args, _kernel=kernel, **kwargs): + result = _kernel(*args, **kwargs) + if not features._active_target_ids(env): + return result + context = env.engine.context + if mode == "mla": + assert context.cache_mla_latent and args[1].shape[-1] == 576 + assert args[4] == 512 and context.is_decode_only() + else: + assert all(mode.upper() + "Tensor" in type(a).__name__ for a in args[:2]) + assert context.padded_active_token_count == 128 + assert any(a.size(0) == 128 for a in args[:2]) + phase = "decode" if context.is_decode_only() else "prefill" + runtime[f"{mode}-{phase}"] += 1 + return result + + monkeypatch.setattr(module, symbol, observed) + + monkeypatch.setattr(features, "_instrument_scenario_runtime", instrument_native) + required = (f"{mode}-decode",) if mode == "mla" else (f"{mode}-prefill", f"{mode}-decode") + await features.test_routed_model_kernels_match_direct(monkeypatch, options, (), required) diff --git a/tests/unit_tests/inference/test_coordinator_batch_invariant_pairwise.py b/tests/unit_tests/inference/test_coordinator_batch_invariant_pairwise.py new file mode 100644 index 00000000000..70f02eadf09 --- /dev/null +++ b/tests/unit_tests/inference/test_coordinator_batch_invariant_pairwise.py @@ -0,0 +1,82 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import pytest +import torch + +from megatron.core.inference.config import AsyncScheduleMode +from megatron.core.inference.contexts.dynamic_context import DynamicInferenceContext +from megatron.core.transformer.custom_layers import batch_invariant_kernels as bik +from tests.unit_tests.inference import coordinator_pairwise_utils as pairwise +from tests.unit_tests.inference import test_coordinator_features_pairwise as features + + +@pytest.mark.internal +@pytest.mark.asyncio +async def test_routed_batch_invariance_at_padding_boundary(monkeypatch, request): + monkeypatch.setattr(DynamicInferenceContext, "TOKEN_ROUNDER", 64) + monkeypatch.setattr(DynamicInferenceContext, "REQUEST_ROUNDER", 4) + bik.enable_batch_invariant_mode("te_native", collective="ordered") + request.addfinalizer(bik.disable_batch_invariant_mode) + from transformer_engine.pytorch import cpp_extensions + from transformer_engine.pytorch.module import layernorm_linear, linear + + modules = {"linear": linear, "layernorm_linear": layernorm_linear} + assert all(module.general_gemm is cpp_extensions.general_gemm for module in modules.values()) + run_routed = features._exercise_all_owners + records, phase, active = {"direct": [], "routed": []}, "direct", [] + + for name, module in modules.items(): + + def traced(*args, _name=name, **kwargs): + if active: + active[0].append((_name, tuple(args[0].shape), tuple(args[1].shape))) + return cpp_extensions.general_gemm(*args, **kwargs) + + monkeypatch.setattr(module, "general_gemm", traced) + + params = pairwise.greedy_params(num_tokens_to_generate=4, return_prompt_tokens=True) + async with pairwise.routed_model( + monkeypatch, + engine_factory=pairwise.batch_invariant_engine, + context_max_requests=128, + context_block_size_tokens=256, + context_buffer_size_gb=0.125, + materialize_only_last_token_logits=False, + async_sched_mode=AsyncScheduleMode.LEGACY, + ) as harness: + controller = harness.engine.controller + forward = controller._dynamic_step_forward_logits + + def observed(tokens, positions): + context = harness.engine.context + ids = features._active_target_ids(harness) + if not ids: + return forward(tokens, positions) + request_index = context.request_ids.tolist().index(min(ids)) + mapping = context.token_to_request_idx[: context.active_token_count] + rows = (mapping == request_index).nonzero().flatten() + active.append([]) + result = forward(tokens, positions) + calls = active.pop() + assert {call[2][0] for call in calls} == {64 if phase == "direct" else 128} + assert {call[0] for call in calls} == set(modules) + if phase == "routed" and context.num_decode_requests: + assert context.num_decode_requests == len(ids) == 65 + input_tokens = tokens[0, rows].clone() + input_positions = positions[0, rows].clone() + logits = controller._all_logits_cuda[0, rows, : harness.config.vocab_size].clone() + records[phase].append((input_tokens, input_positions, logits)) + return result + + controller._dynamic_step_forward_logits = observed + direct, phase = await harness.direct([4], params), "routed" + + def clear_routed(_): + records["routed"].clear() + + await run_routed(harness, [4], params, direct, install=clear_routed, requests_per_client=65) + positions = [[row[1].item() for row in rows] for rows in records.values()] + assert positions == [[0, 1, 2, 3]] * 2 + for expected, actual in zip(*records.values()): + assert all(torch.equal(left, right) for left, right in zip(expected, actual)) + assert torch.isfinite(expected[2]).all() diff --git a/tests/unit_tests/inference/test_coordinator_execution_pairwise.py b/tests/unit_tests/inference/test_coordinator_execution_pairwise.py new file mode 100644 index 00000000000..2ab3d705f83 --- /dev/null +++ b/tests/unit_tests/inference/test_coordinator_execution_pairwise.py @@ -0,0 +1,140 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +from collections import Counter +from contextlib import asynccontextmanager + +import pytest +import torch + +from megatron.core.inference.config import CudaGraphSizingDistribution +from megatron.core.transformer.cuda_graphs import _CudagraphReplayNode +from tests.unit_tests.inference import test_coordinator_features_pairwise as features +from tests.unit_tests.inference import test_coordinator_schedule_pairwise as schedule +from tests.unit_tests.inference.coordinator_pairwise_utils import greedy_params, routed_model + +pytestmark = [pytest.mark.internal, pytest.mark.asyncio] + + +@pytest.mark.parametrize("all_prefills", [False, True]) +@pytest.mark.parametrize("distribution", list(CudaGraphSizingDistribution)) +async def test_routed_prefill_graph_selection(monkeypatch, all_prefills, distribution): + original = schedule.routed_model + + @asynccontextmanager + async def instrumented(*args, **kwargs): + async with original( + *args, + **kwargs, + hidden_size=128, + num_cuda_graphs=3, + force_build_cuda_graphs=True, + use_cuda_graphs_for_non_decode_steps=True, + cuda_graph_all_prefills=all_prefills, + cuda_graph_max_tokens=16, + cuda_graph_mixed_prefill_count=2, + cuda_graph_sizing_distribution=distribution, + ) as h: + replay, observations = _CudagraphReplayNode.forward, [] + + def observed(ctx, runner, first, *inputs): + c = h.engine.context + ids, padded = features._active_target_ids(h), c.padded_batch_dimensions + counts = (c.active_token_count, c.num_prefill_requests, c.num_decode_requests) + result = replay(ctx, runner, first, *inputs) + if h.engine.use_coordinator and any( + len(h.engine.get_request(rid).prompt_tokens) == 65 for rid in ids + ): + observations.append((padded, counts, inputs[0].shape[1])) + return result + + monkeypatch.setattr(_CudagraphReplayNode, "forward", staticmethod(observed)) + yield h + invalid = any( + dim not in h.engine.context.cuda_graph_batch_dimensions_list + or dim.token_count != physical + for dim, _, physical in observations + ) + is_linear = distribution == CudaGraphSizingDistribution.LINEAR + expected = 12 if is_linear else 8 * (1 + all_prefills) + mixed = any( + n == 7 and p and d and dim.token_count == expected + for dim, (n, p, d), _ in observations + ) + prefill = any( + p and not d and dim.token_count > 16 for dim, (_, p, d), _ in observations + ) + assert await h.sync.all_reduce_max(mixed, prefill, invalid) == (1, int(all_prefills), 0) + + monkeypatch.setattr(schedule, "routed_model", instrumented) + await schedule.test_routed_chunk_partition_beside_decode(monkeypatch) + + +@pytest.mark.parametrize("interval", [20, 3, None]) +@pytest.mark.parametrize("synchronous", [False, True]) +async def test_routed_idle_expert_progress(monkeypatch, interval, synchronous): + params = greedy_params(num_tokens_to_generate=32, return_log_probs=True) + params.skip_prompt_log_probs = True + async with routed_model( + monkeypatch, + expert_model_parallel_size=2 if interval else 1, + use_moe_layer_spec=True, + transformer_impl="inference_optimized", + inference_moe_token_dispatcher_type="nccl", + ) as h: + h.engine.disable_ep_consensus = interval is None + h.engine.ep_consensus_interval = interval or 20 + h.engine.use_synchronous_zmq_collectives = synchronous + assert h.dp_size == 2 and h.engine.pg_collection.ep.size() == (2 if interval else 1) + runtime, spans, cadence = Counter(), {False: set(), True: set()}, [] + span_key = "nccl-token-dispatches" if interval else "model-forwards" + model = h.engine.controller.inference_wrapped_model.model + for module in model.modules(): + features._instrument_nccl_dispatch_runtime(module, runtime) + forward, consensus = model.forward, h.engine._ep_establish_consensus + + def observed(*args, **kwargs): + real, before = bool(features._active_target_ids(h)), runtime[span_key] + result = forward(*args, **kwargs) + if h.engine.use_coordinator: + runtime["model-forwards"] += 1 + spans[real].update(range(before, runtime[span_key])) + return result + + def observed_dummy(dummy=h.engine.controller.dummy_forward): + dummy() + runtime["dummy-forwards"] += int(h.engine.use_coordinator) + + async def observed_consensus(work, *args, **kwargs): + runtime["ep-consensus-calls"] += int(h.engine.use_coordinator) + if work and h.engine._last_ep_consensus[0]: + cadence.append(h.engine._ep_consensus_loop_counter) + return await consensus(work, *args, **kwargs) + + monkeypatch.setattr(model, "forward", observed) + monkeypatch.setattr(h.engine.controller, "dummy_forward", observed_dummy) + monkeypatch.setattr(h.engine, "_ep_establish_consensus", observed_consensus) + direct = await h.direct([4, 5, 6], params) + await h.start() + await h.pause() + runtime.clear() + spans, cadence = {False: set(), True: set()}, [] + [result], owners, local_ids = await h.run_paused_batch([(0, [4, 5, 6], params)]) + await h.pause() + assert local_ids == [(0, 0)] and len(owners) == 1 + assert result["status"] == direct["status"] == "COMPLETED" + assert result["generated_tokens"] == direct["generated_tokens"] + scores = torch.tensor([r["generated_log_probs"] for r in (direct, result)]) + assert scores.shape == (2, 32) and torch.isfinite(scores).all() + records = [None] * h.dp_size + torch.distributed.all_gather_object(records, (spans, runtime, cadence)) + owner = next(i for i in range(h.dp_size) if owners[0] == f"mp-coord-{i}".encode()) + real, idle = records[owner][0][True], records[1 - owner][0] + assert real and idle[False] and not idle[True] and (not interval or real <= idle[False]) + for _, stats, _ in records: + assert interval or not stats["ep-consensus-calls"] + assert interval or stats["dummy-forwards"] > 0 + assert stats["nccl-token-dispatches"] == stats["nccl-token-combines"] >= bool(interval) + assert stats["nccl-combine-before-dispatch"] == stats["nccl-dispatch-inflight"] == 0 + calls = records[owner][2] + assert (bool(calls) and all(c % interval == 0 for c in calls)) if interval else not calls + h.assert_retired() diff --git a/tests/unit_tests/inference/test_coordinator_features_pairwise.py b/tests/unit_tests/inference/test_coordinator_features_pairwise.py new file mode 100644 index 00000000000..2d90354ce7c --- /dev/null +++ b/tests/unit_tests/inference/test_coordinator_features_pairwise.py @@ -0,0 +1,372 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Real coordinator routing crossed with model topology and feature execution.""" + +from collections import Counter +from types import SimpleNamespace + +import pytest +import torch + +from megatron.core.transformer.cuda_graphs import _CudagraphReplayNode +from tests.unit_tests.inference.coordinator_pairwise_utils import greedy_params, routed_model +from tests.unit_tests.inference.engines.test_dynamic_engine_async_sched import ( + _instrument_nccl_dispatch_runtime, + _instrument_parallel_runtime, + _instrument_scenario_runtime, +) + +pytestmark = [pytest.mark.internal, pytest.mark.asyncio] + + +def _active_target_ids(harness): + context = harness.engine.context + if getattr(context, "_bookkeeping_no_real_work", False): + return () + ids = context.request_ids[context.paused_request_count : context.total_request_count].tolist() + return tuple(request_id for request_id in ids if request_id in harness.engine.requests) + + +def _install_target_attribution(harness, runtime): + """Attribute completed runtime deltas only to a real coordinator request.""" + target = Counter() + model = harness.engine.controller.inference_wrapped_model.model + forward = model.forward + + def attributed_forward(*args, **kwargs): + ids = _active_target_ids(harness) + before = runtime.copy() + result = forward(*args, **kwargs) + if ids: + target["model-forward"] += 1 + target["prefill-forward"] += int(harness.engine.context.num_prefill_requests > 0) + target["decode-forward"] += int(harness.engine.context.num_decode_requests > 0) + for request_id in ids: + target[("request", request_id)] += 1 + for key, count in runtime.items(): + target[key] += count - before[key] + return result + + model.forward = attributed_forward + return target + + +def _assert_every_owner_executed(harness, target, requests_per_client=1): + target_ids = {key[1] for key in target if isinstance(key, tuple) and key[0] == "request"} + groups = harness.engine.pg_collection + record = { + "dp": torch.distributed.get_rank(groups.dp), + "mp": torch.distributed.get_rank(groups.mp), + "tp": torch.distributed.get_rank(groups.tp), + "pp": torch.distributed.get_rank(groups.pp), + "ids": sorted(target_ids), + "steps": [target[key] for key in ("model-forward", "prefill-forward", "decode-forward")], + } + records = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(records, record) + assert all( + min(item["steps"]) > 0 and len(item["ids"]) == requests_per_client for item in records + ) + assert {item["dp"] for item in records} == set(range(harness.dp_size)) + routed_ids = set() + for owner in range(harness.dp_size): + owner_records = [item for item in records if item["dp"] == owner] + assert {item["mp"] for item in owner_records} == set(range(len(owner_records))) + assert {(item["tp"], item["pp"]) for item in owner_records} == { + (tp, pp) + for tp in range(harness.config.tensor_model_parallel_size) + for pp in range(harness.config.pipeline_model_parallel_size) + } + owner_id_sets = {tuple(item["ids"]) for item in owner_records} + assert len(owner_id_sets) == 1 + owner_ids = set(owner_id_sets.pop()) + assert routed_ids.isdisjoint(owner_ids) + routed_ids.update(owner_ids) + assert routed_ids == set(range(harness.dp_size * requests_per_client)) + + +async def _exercise_all_owners( + harness, prompt, params, direct, runtime=None, install=None, requests_per_client=1 +): + runtime = Counter() if runtime is None else runtime + await harness.start() + target = _install_target_attribution(harness, runtime) + if install is not None: + install(target) + expected_ids = [ + (client, local_id) + for client in range(harness.dp_size) + for local_id in range(requests_per_client) + ] + outputs, owners, local_ids = await harness.run_paused_batch( + [(client, prompt, params) for client, _ in expected_ids] + ) + assert local_ids == expected_ids + assert sorted(Counter(owners.values()).values()) == [requests_per_client] * harness.dp_size + await harness.pause() + for output in outputs: + for key in ("status", "generated_tokens", "prompt_tokens", "generated_text"): + assert output[key] == direct[key], key + _assert_every_owner_executed(harness, target, requests_per_client) + harness.assert_retired() + return target + + +@pytest.mark.parametrize( + "overrides", + [ + pytest.param({"tensor_model_parallel_size": 2}, id="tp2-dp4"), + pytest.param({"pipeline_model_parallel_size": 2}, id="pp2-dp4"), + pytest.param( + { + "tensor_model_parallel_size": 2, + "pipeline_model_parallel_size": 2, + "sequence_parallel": True, + }, + id="tp2-pp2-sp-dp2", + ), + ], +) +async def test_routed_topology_matches_direct_on_every_owner(monkeypatch, overrides): + """TP, PP and the full eight-rank TP2xPP2xDP2 layout execute routed targets.""" + prompt = list(range(4, 16)) + params = greedy_params(num_tokens_to_generate=4, return_prompt_tokens=True) + runtime = Counter() + async with routed_model(monkeypatch, transformer_impl="local", **overrides) as harness: + direct = await harness.direct(prompt, params) + + def install(_target): + if harness.config.sequence_parallel: + from megatron.core.tensor_parallel import layers as tp_layers + + all_gather = tp_layers.dist_all_gather_func + + def traced_all_gather(output, input_, *args, **kwargs): + result = all_gather(output, input_, *args, **kwargs) + runtime["tp-collective:gather_from_sequence_parallel_region"] += 1 + runtime["tp-sp-gather-dimensions"] += int( + output.shape[0] + == input_.shape[0] * harness.config.tensor_model_parallel_size + ) + return result + + monkeypatch.setattr(tp_layers, "dist_all_gather_func", traced_all_gather) + for module in harness.engine.controller.inference_wrapped_model.model.modules(): + if "ParallelLinear" in type(module).__name__: + _instrument_parallel_runtime(module, runtime) + + target = await _exercise_all_owners( + harness, prompt, params, direct, runtime=runtime, install=install + ) + if harness.config.tensor_model_parallel_size > 1: + assert runtime["tp-column-partitions-installed"] > 0 + assert runtime["tp-row-partitions-installed"] > 0 + if harness.config.sequence_parallel: + assert target["tp-collective:gather_from_sequence_parallel_region"] > 0 + assert target["tp-collective:reduce_scatter_to_sequence_parallel_region"] > 0 + assert target["tp-sp-gather-dimensions"] > 0 + assert target["tp-sp-reduce-scatter-dimensions"] > 0 + else: + assert target["tp-collective:reduce_from_tensor_model_parallel_region"] > 0 + + +async def test_routed_cuda_graph_replays_for_every_target(monkeypatch): + """Completed routed decode forwards invoke the production CUDA replay node.""" + prompt = list(range(4, 16)) + params = greedy_params(num_tokens_to_generate=4, return_prompt_tokens=True) + async with routed_model( + monkeypatch, + num_cuda_graphs=4, + force_build_cuda_graphs=True, + use_cuda_graphs_for_non_decode_steps=False, + cuda_graph_max_tokens=16, + ) as harness: + direct = await harness.direct(prompt, params) + + def install(target): + replay = _CudagraphReplayNode.forward + + def attributed_replay(*args, **kwargs): + result = replay(*args, **kwargs) + if _active_target_ids(harness): + target.update(("model-forward", "decode-forward", "cuda-graph-replay")) + return result + + monkeypatch.setattr(_CudagraphReplayNode, "forward", staticmethod(attributed_replay)) + + target = await _exercise_all_owners(harness, prompt, params, direct, install=install) + assert target["cuda-graph-replay"] > 0 + + +async def test_routed_mtp_executes_inner_steps_on_every_target(monkeypatch): + """Coordinator-owned speculative decoding calls the real inner MTP model step.""" + prompt = list(range(4, 16)) + params = greedy_params(num_tokens_to_generate=6, return_prompt_tokens=True) + async with routed_model(monkeypatch, num_speculative_tokens=1) as harness: + direct = await harness.direct(prompt, params) + + def install(target): + model = harness.engine.controller.inference_wrapped_model.model + compute = model.compute_mtp_single_step + + def traced_compute(*args, **kwargs): + result = compute(*args, **kwargs) + if _active_target_ids(harness): + target["mtp-inner-step"] += 1 + return result + + model.compute_mtp_single_step = traced_compute + + target = await _exercise_all_owners(harness, prompt, params, direct, install=install) + assert target["mtp-inner-step"] > 0 + assert harness.engine._spec_steps > 0 + assert int(harness.engine._spec_tokens_proposed_per_pos.sum()) > 0 + + +@pytest.mark.parametrize("mixer", ["mamba", "gdp", "gdn"]) +async def test_routed_hybrid_updates_owned_recurrent_state(monkeypatch, mixer): + """Completed target decode updates the target's mapped recurrent state slot.""" + prompt = list(range(4, 16)) + params = greedy_params(num_tokens_to_generate=4, return_prompt_tokens=True) + async with routed_model(monkeypatch, model_provider="hybrid", ssm_mixer=mixer) as harness: + direct = await harness.direct(prompt, params) + + def install(target): + model = harness.engine.controller.inference_wrapped_model.model + context = harness.engine.context + expected = { + "mamba": "MambaMixer", + "gdp": "GatedDeltaProductMixer", + "gdn": "GatedDeltaNet", + } + selected = [ + module for module in model.modules() if type(module).__name__ == expected[mixer] + ] + assert selected, expected[mixer] + for module in selected: + module.register_forward_hook( + lambda *_: target.update( + {"mixer-forward": int(bool(_active_target_ids(harness)))} + ) + ) + forward = model.forward + + def traced_forward(*args, **kwargs): + snapshots = [] + if context.is_decode_only(): + for request_id in _active_target_ids(harness): + row = int( + ( + context.request_ids[: context.total_request_count] == request_id + ).nonzero()[0] + ) + slot = int(context.mamba_metadata.request_to_mamba_state_idx[row]) + assert 0 <= slot < context.max_requests + snapshots.append( + ( + slot, + context.mamba_conv_states[:, slot].clone(), + context.mamba_ssm_states[:, slot].clone(), + ) + ) + result = forward(*args, **kwargs) + for slot, conv, ssm in snapshots: + assert not torch.equal(conv, context.mamba_conv_states[:, slot]) + assert not torch.equal(ssm, context.mamba_ssm_states[:, slot]) + target["owned-state-updates"] += 1 + return result + + model.forward = traced_forward + + target = await _exercise_all_owners(harness, prompt, params, direct, install=install) + assert harness.engine.context.is_hybrid_model + assert target["mixer-forward"] > 0 + assert target["owned-state-updates"] > 0 + + +@pytest.mark.parametrize( + "options,signals,required", + [ + ({"fp8": True, "hidden_size": 128}, ("fp8",), ("fp8-quantized-forwards",)), + ( + { + # FI 0.6.6 NeoX D16 corrupts adjacent heads; exercise native D64. + "hidden_size": 256, + "position_embedding_type": "rope", + "use_flashinfer_fused_rope": True, + }, + ("fused-rope",), + ("fused-rope-kernel",), + ), + ( + {"window_size": (4, 0), "softmax_type": "off-by-one"}, + ("softmax-sink",), + ("swa-kernel-calls", "sink-correction-calls"), + ), + ( + {"window_size": (4, 0), "window_attn_skip_freq": 2, "softmax_type": "learnable"}, + ("softmax-sink",), + ("swa-kernel-calls", "full-attention-kernel-calls", "sink-correction-calls"), + ), + ], + ids=["fp8", "fused-rope", "swa-sink", "alternating-swa-sink"], +) +async def test_routed_model_kernels_match_direct(monkeypatch, options, signals, required): + prompt = list(range(4, 16)) + params = greedy_params(num_tokens_to_generate=4, return_prompt_tokens=True) + async with routed_model(monkeypatch, **options) as harness: + missing = torch.tensor(False, device="cuda") + finite = [] + sample_kernel = (sampling := harness.engine.controller._sampling).sample_kernel + reduce = torch.distributed._functional_collectives.all_reduce + + def observed_sample(logits, n, context, **kwargs): + if _active_target_ids(harness): + assert n == 1 + gather_indices = kwargs.get("gather_indices") + rows = logits[gather_indices[:n]] if gather_indices is not None else logits[:n] + finite.append(torch.isfinite(rows[:, : harness.config.vocab_size]).all()) + return sample_kernel(logits, n, context, **kwargs) + + sampling.sample_kernel = observed_sample + direct = await harness.direct(prompt, params) + assert reduce(torch.stack(finite or [missing]).all(), "min", torch.distributed.group.WORLD) + finite.clear() + runtime = Counter() + _instrument_scenario_runtime( + SimpleNamespace(engine=harness.engine), SimpleNamespace(signals=signals), runtime + ) + target = await _exercise_all_owners(harness, prompt, params, direct, runtime=runtime) + assert reduce(torch.stack(finite or [missing]).all(), "min", torch.distributed.group.WORLD) + for counter in required: + assert target[counter] > 0, counter + + +async def test_routed_moe_participates_in_ep_collectives(monkeypatch): + """Every coordinator target dispatches and combines through its EP collective.""" + prompt = list(range(4, 16)) + params = greedy_params(num_tokens_to_generate=4, return_prompt_tokens=True) + runtime = Counter() + async with routed_model( + monkeypatch, + expert_model_parallel_size=2, + use_moe_layer_spec=True, + inference_moe_token_dispatcher_type="nccl", + transformer_impl="inference_optimized", + ) as harness: + direct = await harness.direct(prompt, params) + + def install(_target): + for module in harness.engine.controller.inference_wrapped_model.model.modules(): + if type(module).__name__ == "MoELayer": + _instrument_nccl_dispatch_runtime(module, runtime) + + target = await _exercise_all_owners( + harness, prompt, params, direct, runtime=runtime, install=install + ) + assert torch.distributed.get_world_size(harness.engine.pg_collection.ep) == 2 + assert runtime["nccl-dispatchers-installed"] > 0 + assert target["nccl-token-dispatches"] > 0 + assert target["nccl-token-dispatches"] == target["nccl-token-combines"] + assert target["nccl-combine-before-dispatch"] == 0 + assert target["nccl-dispatch-inflight"] == 0 diff --git a/tests/unit_tests/inference/test_coordinator_lifecycle_pairwise.py b/tests/unit_tests/inference/test_coordinator_lifecycle_pairwise.py new file mode 100644 index 00000000000..3a313b6e855 --- /dev/null +++ b/tests/unit_tests/inference/test_coordinator_lifecycle_pairwise.py @@ -0,0 +1,284 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import asyncio +import copy +from unittest import mock + +import pytest +import torch + +from megatron.core.inference import unified_memory as unified_memory_module +from megatron.core.inference.config import AsyncScheduleMode, KVCacheManagementMode +from megatron.core.inference.contexts import dynamic_context as dynamic_context_module +from megatron.core.inference.engines.dynamic_engine import EngineState +from megatron.core.inference.headers import Headers +from megatron.core.transformer.cuda_graphs import _CudaGraphRunner +from tests.unit_tests.inference.coordinator_pairwise_utils import greedy_params, routed_model, until +from tests.unit_tests.inference.engines import test_dynamic_engine as engine_tests +from tests.unit_tests.inference.test_coordinator_features_pairwise import _exercise_all_owners + +pytestmark = [pytest.mark.internal, pytest.mark.asyncio] + + +async def _pause_active(h, monkeypatch, prompt, params): + arrived, release = asyncio.Event(), asyncio.Event() + step = h.engine.async_step + + async def gated_step(): + if arrived.is_set(): + return None + result = await step() + if h.rank == 0 and not arrived.is_set() and 0 in h.engine.requests: + if len(h.engine.get_request(0).generated_tokens) >= 2: + arrived.set() + await asyncio.wait_for(release.wait(), timeout=30) + return result + + monkeypatch.setattr(h.engine, "async_step", gated_step) + future = None + if h.rank == 0: + future = h.clients[0].add_request(prompt, copy.deepcopy(params)) + await asyncio.wait_for(arrived.wait(), timeout=30) + h.clients[0].pause_engines() + await until(lambda: any(e["header"] == Headers.PAUSE for e in h.service.events)) + release.set() + await asyncio.wait_for(h.engine.wait_until(EngineState.PAUSED), timeout=60) + monkeypatch.setattr(h.engine, "async_step", step) + await h.barrier() + if h.rank == 0: + assert 0 < len(h.engine.get_request(0).generated_tokens) < params.num_tokens_to_generate + assert h.witnesses and not future.done() + return future + + +@pytest.mark.parametrize("mode", list(AsyncScheduleMode)) +@pytest.mark.parametrize("residency", list(KVCacheManagementMode)) +async def test_routed_suspend_resumes_active_request(monkeypatch, mode, residency): + """Supported residency transitions preserve a partly generated routed result.""" + prompt, params = list(range(4, 20)), greedy_params(num_tokens_to_generate=16) + offload = residency == KVCacheManagementMode.OFFLOAD + if offload: + saver = dynamic_context_module.torch_memory_saver + for name in ("region", "pause", "resume"): + monkeypatch.setattr(saver, name, mock.Mock(wraps=getattr(saver, name))) + async with routed_model( + monkeypatch, + async_sched_mode=mode, + kv_cache_management_mode=residency.value, + static_kv_memory_pointers=offload, + num_cuda_graphs=1 if offload else None, + force_build_cuda_graphs=offload, + use_cuda_graphs_for_non_decode_steps=not offload, + ) as h: + direct = await h.direct(prompt, params) + await h.start() + context, replays = h.engine.context, [] + if offload: + replay = _CudaGraphRunner.replay_graph_capture + + def observed_replay(runner, *args, **kwargs): + result = replay(runner, *args, **kwargs) + active = slice(context.paused_request_count, context.total_request_count) + ids = context.request_ids[active] + live = 0 in h.engine.requests and not context._bookkeeping_no_real_work + replays.extend([True] if live and 0 in ids else []) + return result + + monkeypatch.setattr(_CudaGraphRunner, "replay_graph_capture", observed_replay) + future = await _pause_active(h, monkeypatch, prompt, params) + pointer = context.memory_buffer.data_ptr() + if h.rank == 0: + route = dict(h.service.coordinator.request_id_to_rank) + assert route[0] == f"mp-coord-{h.rank}".encode() and 0 in h.engine.requests + row = context.request_ids.tolist().index(0) + blocks = context.request_to_kv_block_ids[row] + blocks = blocks[: context.request_kv_block_counts[row]].long() + target_kv = context.memory_buffer[:, :, blocks].view(torch.int16).clone() + before = list(h.engine.get_request(0).generated_tokens) + h.clients[0].suspend_engines() + await asyncio.wait_for(h.engine.wait_until(EngineState.SUSPENDED), timeout=60) + assert not context.is_tensor_state_allocated + if residency == KVCacheManagementMode.PERSIST: + assert context.memory_buffer.data_ptr() == pointer + elif offload: + saver.region.assert_called_once_with(tag=context.TMS_TAG, enable_cpu_backup=True) + saver.pause.assert_called_once_with(context.TMS_TAG) + else: + assert not hasattr(context, "memory_buffer") + assert not h.engine.controller._async_sched_logits.is_valid + await h.barrier() + if h.rank == 0: + assert h.service.coordinator.request_id_to_rank == route and not future.done() + h.clients[0].resume_engines() + await asyncio.wait_for(h.engine.wait_until(EngineState.RESUMED), timeout=60) + assert context.is_tensor_state_allocated and h.engine.state == EngineState.PAUSED + assert h.engine._weight_epoch == 1 + if offload: + assert context.memory_buffer.data_ptr() == pointer + saver.resume.assert_called_once_with(context.TMS_TAG) + if h.rank == 0: + assert context.memory_buffer[:, :, blocks].view(torch.int16).equal(target_kv) + replays[:] = h.witnesses[:] = [] + await h.barrier() + await h.unpause() + if h.rank == 0: + final = await asyncio.wait_for(future, timeout=60) + assert final["request_id"] == 0 and final["generated_tokens"][: len(before)] == before + assert final["generated_tokens"] == direct["generated_tokens"] + assert final["status"] == direct["status"] + if residency == KVCacheManagementMode.RECOMPUTE: + assert sum(bool(s["prefill"]) for s in h.witnesses) >= 2 + if offload: + assert await h.sync.all_reduce_max(int(bool(replays))) == 1 + await h.barrier() + h.assert_retired() + + +@pytest.mark.parametrize("active", [False, True], ids=["queued", "active"]) +async def test_routed_abort_retires_only_cancelled_request(monkeypatch, active): + prompt, params = list(range(4, 20)), greedy_params(num_tokens_to_generate=16) + async with routed_model(monkeypatch) as h: + direct = await h.direct(prompt, params) + await h.start() + if active: + cancelled = await _pause_active(h, monkeypatch, prompt, params) + else: + await h.pause() + if h.rank == 0: + cancelled = h.clients[0].add_request(prompt, copy.deepcopy(params)) + await until(lambda: 0 in h.engine.requests) + assert not h.witnesses + if h.rank == 0: + survivor_id, survivor = h.clients[1].add_request_with_id(prompt, copy.deepcopy(params)) + assert survivor_id == 0 + await until(lambda: len(h.service.coordinator.request_id_to_rank) == 2) + h.clients[0].abort_request(0) + assert cancelled.cancelled() + await until(lambda: any(e["header"] == Headers.ABORT_REQUEST for e in h.service.events)) + await h.barrier() + await h.unpause() + if h.rank == 0: + final = await asyncio.wait_for(survivor, timeout=60) + assert final["generated_tokens"] == direct["generated_tokens"] + await until(lambda: not h.service.coordinator.request_id_to_rank) + await until(lambda: not h.clients[0].aborted_request_ids) + await h.barrier() + h.assert_retired() + + +@pytest.mark.parametrize("synchronous", [False, True]) +async def test_routed_epoch_change_reaches_all_engines(monkeypatch, synchronous): + """A paused generation update reaches idle peers and the active request.""" + prompt, params = list(range(4, 20)), greedy_params(num_tokens_to_generate=16) + async with routed_model(monkeypatch) as h: + direct = await h.direct(prompt, params) + await h.start() + h.engine.use_synchronous_zmq_collectives = synchronous + calls, operation = [], h.engine.world_zmq_communicator.all_reduce_max + + async def observed_reduce(*args, **kwargs): + result = await operation(*args, **kwargs) + calls.append(kwargs["async_op"]) + return result + + monkeypatch.setattr(h.engine.world_zmq_communicator, "all_reduce_max", observed_reduce) + future = await _pause_active(h, monkeypatch, prompt, params) + if h.rank == 0: + h.clients[0].set_generation_epoch(7) + await until(lambda: h.engine._generation_epoch == 7) + if h.rank == 0: + request = h.engine.get_request(0) + assert (request.policy_epoch, request.kv_cache_epoch) == ([(0, 7)], [(0, 7)]) + await h.barrier() + await h.unpause() + assert len(calls) >= 2 and set(calls) == {not synchronous} + if h.rank == 0: + final = await asyncio.wait_for(future, timeout=60) + assert final["generated_tokens"] == direct["generated_tokens"] + assert (final["policy_epoch"], final["kv_cache_epoch"]) == ([[0, 7]], [[0, 7]]) + await h.barrier() + h.assert_retired() + + +async def test_routed_allocator_pressure_preserves_victim_identity(monkeypatch): + prompt, params = list(range(4, 20)), greedy_params() + async with routed_model( + monkeypatch, context_buffer_size_gb=0.00004, context_paused_buffer_size_gb=0.000008 + ) as h: + direct = await h.direct(prompt, params) + context, transitions = h.engine.context, [] + allocator, update = context.kv_block_allocator, context.update_requests + resume_spy = mock.Mock(wraps=context.resume_paused_requests) + assert (allocator.pool_size, allocator.paused_limit) == (5, 1) + + def observed_update(*args, **kwargs): + result = update(*args, **kwargs) + evicted = result["evict_request_ids"] + if evicted is not None: + paused = result["newly_paused_request_ids"] + transitions.append((paused.shape, paused.flatten().tolist(), evicted.tolist())) + return result + + monkeypatch.setattr(context, "resume_paused_requests", resume_spy) + monkeypatch.setattr(context, "update_requests", observed_update) + await h.start() + count = h.dp_size * (allocator.pool_size - 2) + finals, owners, local_ids = await h.run_paused_batch([(0, prompt, params)] * count) + reference = (direct["generated_tokens"], direct["status"]) + assert [(r["generated_tokens"], r["status"]) for r in finals] == [reference] * count + assert local_ids == [(0, request_id) for request_id in range(count)] + [(shape, paused, evicted)] = transitions + assert resume_spy.call_args_list[0].args[1].shape == (allocator.pool_size - 2,) + assert paused == evicted and shape == (1,) + target = evicted[0] + assert owners[target] == f"mp-coord-{h.rank}".encode() + assert sum(bool(s["prefill"]) for s in h.witnesses if target in s["ids"]) >= 2 + h.assert_retired() + + +async def test_routed_uvm_drained_reset_preserves_live_controls(monkeypatch): + config_constructor = engine_tests.InferenceConfig.__init__ + pool_spy = mock.Mock(wraps=unified_memory_module.MemPool) + + def uvm_config(config, *args, **kwargs): + kwargs["unified_memory_level"] = 1 + config_constructor(config, *args, **kwargs) + + def engine_factory(config): + with mock.patch.object(engine_tests.InferenceConfig, "__init__", uvm_config): + return engine_tests.DynamicInferenceEngineTestBase._build_test_env(config).engine + + monkeypatch.setattr(unified_memory_module, "MemPool", pool_spy) + prompt, params = list(range(4, 20)), greedy_params() + async with routed_model(monkeypatch, engine_factory=engine_factory) as h: + context = h.engine.context + assert context.unified_memory_level == pool_spy.call_count == 1 + assert pool_spy.call_args.kwargs["allocator"] is unified_memory_module._alloc + pointer, end = context.memory_buffer.data_ptr(), context.memory_buffer.nbytes + assert any( + s["address"] <= pointer and pointer + end <= s["address"] + s["total_size"] + for s in context.unified_memory_mempool.snapshot() + ) + unified_memory_module.advise_managed_tensor_preferred_location( + context.memory_buffer, device=-1 + ) + unified_memory_module.prefetch_managed_tensor(context.memory_buffer, device=-1) + unified_memory_module.prefetch_managed_tensor(context.memory_buffer, device=h.rank) + torch.cuda.synchronize() + direct = await h.direct(prompt, params) + await _exercise_all_owners(h, prompt, params, direct) + names = ("_cond", "_state_events", "_pending_signals", "world_zmq_communicator") + runtime = {name: getattr(h.engine, name) for name in names} + h.engine.reset() + assert all(getattr(h.engine, name) is value for name, value in runtime.items()) + assert h.engine.use_coordinator and h.engine.state == EngineState.PAUSED + assert h.engine._state_events[EngineState.PAUSED].is_set() + h.witnesses.clear() + reference = direct["generated_tokens"], direct["status"] + await h.barrier() + await h.unpause() + if h.rank == 0: + final = await asyncio.wait_for(h.clients[0].add_request(prompt, params), 60) + assert (final["generated_tokens"], final["status"]) == reference + await h.barrier() + assert await h.sync.all_reduce_max(bool(h.witnesses)) == 1 + h.assert_retired() diff --git a/tests/unit_tests/inference/test_coordinator_media_pairwise.py b/tests/unit_tests/inference/test_coordinator_media_pairwise.py new file mode 100644 index 00000000000..ff4becbbc58 --- /dev/null +++ b/tests/unit_tests/inference/test_coordinator_media_pairwise.py @@ -0,0 +1,309 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import asyncio +import copy + +import pytest +import torch +from PIL import Image + +from megatron.core.inference.config import ( + ImageProcessingConfig, + MediaCacheCoordinatorPolicy, + PrefixCachingCoordinatorPolicy, + PrefixCachingEvictionPolicy, +) +from megatron.core.inference.engines.dynamic_engine import DynamicInferenceEngine, EngineState +from megatron.core.inference.headers import Headers +from megatron.core.inference.model_inference_wrappers.multimodal.vlm_inference_wrapper import ( + VLMInferenceWrapper, +) +from megatron.core.inference.text_generation_controllers.text_generation_controller import ( + TextGenerationController, +) +from megatron.core.inference.text_generation_server.dynamic_text_gen_server.image_preprocessing import ( + preprocess_image, +) +from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_layer_with_transformer_engine_submodules, +) +from megatron.core.models.multimodal.llava_model import LLaVAModel +from megatron.core.transformer.spec_utils import ModuleSpec, get_submodules +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.transformer.transformer_layer import TransformerLayer +from tests.unit_tests.inference.coordinator_pairwise_utils import greedy_params, routed_model + + +def _install_vlm(h, monkeypatch, dynamic, allow_stale): + language = h.engine.controller.model_config + language.language_model_type = "dummy" + config = dict( + num_layers=1, + hidden_size=16, + num_attention_heads=2, + params_dtype=torch.bfloat16, + use_cpu_initialization=True, + ) + vision = TransformerConfig(**config) + vision.vision_model_type = "radio" if dynamic else "clip" + projection = TransformerConfig( + **{ + **config, + "hidden_size": language.hidden_size, + "ffn_hidden_size": 32, + "num_attention_heads": 1, + } + ) + submodules = get_gpt_layer_with_transformer_engine_submodules() + layer = ModuleSpec(module=TransformerLayer, submodules=submodules) + torch.manual_seed(887) + model = ( + LLaVAModel( + language_transformer_config=language, + language_transformer_layer_spec=layer, + language_vocab_size=h.config.vocab_size, + language_max_sequence_length=h.config.max_sequence_length, + vision_transformer_config=vision, + vision_transformer_layer_spec=copy.deepcopy(layer), + drop_vision_class_token=True, + vision_projection_config=projection, + vision_projection_layer_spec=copy.deepcopy(get_submodules(submodules.mlp)), + img_h=32, + img_w=32, + patch_dim=16, + dynamic_resolution=dynamic, + class_token_len=0 if dynamic else None, + radio_disable_cpe=dynamic, + ) + .cuda() + .to(torch.bfloat16) + .eval() + ) + context = h.engine.context + context.config.vision_embedding_cache_max_bytes = 1024 * 1024 + context.config.allow_stale_multimodal_embeddings = allow_stale + wrapper = VLMInferenceWrapper(model, context) + wrapper.prep_model_for_inference() + monkeypatch.setattr(h.tokenizer, "convert_tokens_to_ids", lambda token: 99, raising=False) + h.engine = DynamicInferenceEngine(TextGenerationController(wrapper, h.tokenizer), context) + admissions, calls = [], {"vision": 0, "projection": 0} + for name, module in (("vision", model.vision_model), ("projection", model.vision_projection)): + forward = module.forward + + def observed(*args, forward=forward, name=name, **kwargs): + result = forward(*args, **kwargs) + calls[name] += 1 + return result + + monkeypatch.setattr(module, "forward", observed) + build = h.engine._build_vlm_request + + def build_request(**kwargs): + before = dict(calls) + request = build(**kwargs) + delta = tuple(calls[name] - before[name] for name in calls) + admissions.append( + ( + request.request_id, + h.rank, + request.block_hash_salt, + *delta, + h.engine._weight_epoch, + request.image_embeddings.detach().cpu(), + ) + ) + return request + + monkeypatch.setattr(h.engine, "_build_vlm_request", build_request) + + def observe_language(): + forward = model.language_model.forward + + def observed(*args, **kwargs): + active = context.request_ids[ + context.paused_request_count : context.total_request_count + ].tolist() + ids = [rid for rid in active if rid in h.engine.requests] + mask = context.current_image_token_mask() if context.has_vlm_data else None + positions = 0 if mask is None else int((mask >= 0).sum().item()) + injected = False + if positions: + assert len(ids) == 1, "The schedule must have one live media target per owner" + image_positions = mask[0] >= 0 + actual = kwargs["decoder_input"][image_positions, 0] + embeddings = context.current_image_embeddings() + expected = embeddings[mask[0, image_positions], 0].to(actual.dtype) + injected = torch.equal(actual, expected) + cached = {rid: h.engine.get_request(rid).num_cached_tokens for rid in ids} + result = forward(*args, **kwargs) + if ids: + h.witnesses.append((h.rank, ids, positions, injected, cached)) + return result + + monkeypatch.setattr(model.language_model, "forward", observed) + + monkeypatch.setattr(h, "_observe_forwards", observe_language) + return admissions + + +@pytest.mark.internal +@pytest.mark.asyncio +@pytest.mark.parametrize("policy", list(MediaCacheCoordinatorPolicy)) +@pytest.mark.parametrize( + "dynamic", [False, True, "video"], ids=["clip", "radio-dynamic", "radio-video"] +) +@pytest.mark.parametrize("allow_stale", [False, True], ids=["invalidate", "retain"]) +async def test_routed_media_affinity_uses_real_cached_embeddings( + monkeypatch, policy, dynamic, allow_stale +): + prompt = [*range(4, 20), 99, 20, 21] + params = greedy_params(return_prompt_tokens=True) + modality = "video" if dynamic == "video" else "image" + if dynamic: + pixels = torch.arange(32 * 48 * 3).remainder(251).byte().reshape(32, 48, 3) + image_config = ImageProcessingConfig( + patch_dim=16, + dynamic_resolution=True, + dynamic_resolution_max_patches=6, + pixel_mean=[0.0] * 3, + pixel_std=[1.0] * 3, + ) + + def raw_media(values): + imgs, imgs_sizes = preprocess_image(Image.fromarray(values.numpy()), image_config) + media = {"imgs": imgs, "imgs_sizes": imgs_sizes} + if dynamic == "video": + media.update( + imgs=torch.cat((imgs, imgs.flip(1)), dim=1), + imgs_sizes=imgs_sizes.repeat(2, 1), + num_frames=torch.tensor([2]), + ) + return media + + image_a = raw_media(pixels) + image_b = raw_media(pixels.flip(1).contiguous()) + else: + image_a = dict( + imgs=torch.arange(3 * 32 * 32, dtype=torch.float32).reshape(1, 3, 32, 32) / 3072, + num_tiles=torch.tensor([1]), + num_img_embeddings_per_tile=4, + ) + image_b = {**image_a, "imgs": image_a["imgs"].flip(-1).contiguous()} + async with routed_model( + monkeypatch, + enable_prefix_caching=True, + prefix_caching_eviction_policy=PrefixCachingEvictionPolicy.LRU, + coordinator_options={ + "prefix_caching_coordinator_policy": PrefixCachingCoordinatorPolicy.LOAD_BALANCED, + "media_cache_coordinator_policy": policy, + "vision_embedding_cache_enabled": True, + }, + ) as h: + assert h.dp_size == 2, "This owner-selection schedule is a two-replica row" + admissions = _install_vlm(h, monkeypatch, bool(dynamic), allow_stale) + references = [] + for media in (image_a, image_b): + future = h.engine.add_request(10001, prompt, copy.deepcopy(params), **media) + while h.engine.has_unfinished_requests(): + await h.engine.async_step() + references.append(list((await future).merge().generated_tokens)) + h.engine.reset() + admissions.clear() + await h.start() + submissions = [(0, [4, 5, 6], params), (1, prompt, params, {modality: image_a})] + outputs, _, local_ids = await h.run_paused_batch(submissions) + assert local_ids == [(0, 0), (1, 0)] + assert outputs[1]["generated_tokens"] == references[0] + cache_before = { + key: embedding.detach().cpu().clone() + for key, embedding in h.engine._vision_embedding_cache.items() + } + assert set(cache_before) == {row[2] for row in admissions} + await h.pause() + if h.rank == 0: + h.clients[0].suspend_engines() + await asyncio.wait_for(h.engine.wait_until(EngineState.SUSPENDED), timeout=60) + assert h.engine._weight_epoch == 0 + if allow_stale: + assert set(h.engine._vision_embedding_cache) == set(cache_before) + assert all( + torch.equal(h.engine._vision_embedding_cache[key].cpu(), value) + for key, value in cache_before.items() + ) + else: + assert not h.engine._vision_embedding_cache + await h.barrier() + if h.rank == 0: + h.clients[0].resume_engines() + await asyncio.wait_for(h.engine.wait_until(EngineState.RESUMED), timeout=60) + assert h.engine._weight_epoch == 1 and h.engine.state == EngineState.PAUSED + await h.barrier() + await h.unpause() + if h.rank == 0: + for media, expected in zip((image_a, image_b), references): + result = await asyncio.wait_for( + h.clients[0].add_request( + prompt, copy.deepcopy(params), multi_modal_data={modality: media} + ), + timeout=60, + ) + assert result["generated_tokens"] == expected + submits = [e for e in h.service.events if e["header"] == Headers.SUBMIT_REQUEST] + assert len(submits) == 4 + assert submits[0]["sender"] == submits[2]["sender"] == submits[3]["sender"] + assert submits[1]["sender"] != submits[0]["sender"] + routes = [] + for event in submits: + new_ids = set(event["after"]) - set(event["before"]) + assert len(new_ids) == 1, event + request_id = new_ids.pop() + routes.append((request_id, event["after"][request_id])) + text_route, warm_route, repeat_route, different_route = routes + coordinator = h.service.coordinator + assert text_route[1] != warm_route[1], "Competing text must put the warm on its peer" + expected_repeat_owner = ( + warm_route[1] if policy == MediaCacheCoordinatorPolicy.AFFINITY else text_route[1] + ) + assert repeat_route[1] == expected_repeat_owner + assert different_route[1] == text_route[1] + route_info = [(rid, coordinator.identity_to_rank_index[owner]) for rid, owner in routes] + await h.barrier() + route_payload = [route_info if h.rank == 0 else None] + torch.distributed.broadcast_object_list(route_payload, src=0) + route_info = route_payload[0] + gathered = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(gathered, (admissions, h.witnesses)) + all_admissions = [row for rank_data in gathered for row in rank_data[0]] + all_witnesses = [row for rank_data in gathered for row in rank_data[1]] + targets = [rid for rid, _ in route_info[1:]] + warm_id, repeat_id, different_id = targets + by_id = {row[0]: row for row in all_admissions} + assert len(all_admissions) == 3 and set(by_id) == set(targets) + assert all( + by_id[rid][1] == owner + and any(step[0] == owner and rid in step[1] for step in all_witnesses) + for rid, owner in route_info[1:] + ) + + expected_calls = ( + 1, + 0 if allow_stale and policy == MediaCacheCoordinatorPolicy.AFFINITY else 1, + 1, + ) + expected_positions = 12 if dynamic == "video" else 6 if dynamic else 4 + for request_id, count in zip(targets, expected_calls): + assert by_id[request_id][3] == by_id[request_id][4] == count + assert any( + request_id in step[1] and step[2] == expected_positions and step[3] + for step in all_witnesses + ), f"Request {request_id} must consume its own projected image embeddings" + + assert by_id[warm_id][5] == 0 + assert by_id[repeat_id][5] == by_id[different_id][5] == 1 + assert by_id[repeat_id][2] == f"w1\0{by_id[warm_id][2]}" + assert by_id[repeat_id][2] != by_id[different_id][2] + assert torch.equal(by_id[warm_id][6], by_id[repeat_id][6]) + assert not torch.equal(by_id[warm_id][6], by_id[different_id][6]) + different_steps = [step for step in all_witnesses if different_id in step[1]] + assert different_steps and all(step[4][different_id] == 0 for step in different_steps) + h.assert_retired() diff --git a/tests/unit_tests/inference/test_coordinator_pairwise.py b/tests/unit_tests/inference/test_coordinator_pairwise.py new file mode 100644 index 00000000000..f6a3865f60f --- /dev/null +++ b/tests/unit_tests/inference/test_coordinator_pairwise.py @@ -0,0 +1,126 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Real DP routing crossed with inference behavior, not dummy-engine plumbing.""" + +import asyncio +import copy + +import pytest +import torch + +from megatron.core.inference.config import ( + AsyncScheduleMode, + PrefixCachingCoordinatorPolicy, + PrefixCachingEvictionPolicy, +) +from megatron.core.inference.headers import Headers +from tests.unit_tests.inference.coordinator_pairwise_utils import greedy_params, routed_model, until + + +async def _collect(stream): + return [item async for item in stream] + + +@pytest.mark.internal +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", list(AsyncScheduleMode)) +@pytest.mark.parametrize("streaming_interval", [0, 1, 3]) +async def test_routed_output_matches_direct(monkeypatch, mode, streaming_interval): + """Two clients both use ID zero; each DP owner runs the direct schedule.""" + prompt = list(range(4, 20)) + params = greedy_params(return_log_probs=True, return_prompt_tokens=True, top_n_logprobs=3) + async with routed_model( + monkeypatch, async_sched_mode=mode, materialize_only_last_token_logits=False + ) as h: + direct = await h.direct(prompt, params) + await h.start() + await h.pause() + pending = [] + if h.rank == 0: + for client in h.clients: + p = copy.deepcopy(params) + if streaming_interval: + p.streaming_interval = streaming_interval + stream = client.add_request_streaming(prompt, p) + assert stream.request_id == 0 + pending.append(asyncio.create_task(_collect(stream))) + else: + request_id, future = client.add_request_with_id(prompt, p) + assert request_id == 0 + pending.append(future) + await until(lambda: len(h.service.coordinator.request_id_to_rank) == h.dp_size) + assert sorted(h.service.coordinator._pending_counts.tolist()) == [1] * h.dp_size + await h.barrier() + await h.unpause() + if h.rank == 0: + outputs = await asyncio.wait_for(asyncio.gather(*pending), timeout=60) + engine_request_ids = [] + for output in outputs: + if streaming_interval: + assert len([item for item in output if "final" in item]) == 1 + assert "final" in output[-1] + final = output[-1]["final"] + partials = [item["partial"] for item in output if "partial" in item] + assert partials, "A final-only stream does not cover partial delivery" + tokens = [token for part in partials for token in part["new_tokens"]] + assert tokens == final["generated_tokens"][: len(tokens)] + assert all(len(p["new_tokens"]) >= streaming_interval for p in partials) + for part in partials: + assert len(part["new_log_probs"]) == len(part["new_tokens"]) + else: + final = output + engine_request_ids.append(final["request_id"]) + for key in ("status", "generated_tokens", "prompt_tokens"): + assert final[key] == direct[key], key + for key in ("prompt_log_probs", "generated_log_probs"): + torch.testing.assert_close(torch.tensor(final[key]), torch.tensor(direct[key])) + assert final["generated_text"] == h.tokenizer.detokenize(final["generated_tokens"]) + assert sorted(engine_request_ids) == list(range(h.dp_size)) + replies = [e for e in h.service.events if e["header"] == Headers.ENGINE_REPLY] + assert sum(len(e["metadata"][1]) for e in replies) == h.dp_size + await h.barrier() + assert h.witnesses, "Every DP replica must execute its assigned real request" + assert {rid for step in h.witnesses for rid in step["ids"]} == {h.rank} + assert any(step["prefill"] for step in h.witnesses) + assert any(step["decode"] for step in h.witnesses) + if mode == AsyncScheduleMode.ASYNC: + async_steps = [step["async_step"] for step in h.witnesses] + assert max(async_steps) > min(async_steps) + h.assert_retired() + + +@pytest.mark.internal +@pytest.mark.asyncio +@pytest.mark.parametrize( + "policy", + [ + PrefixCachingCoordinatorPolicy.LONGEST_PREFIX, + PrefixCachingCoordinatorPolicy.FIRST_PREFIX_BLOCK, + ], +) +async def test_routed_prefix_reuse_executes_on_warm_owner(monkeypatch, policy): + """Affinity must lead to real KV reuse, not just matching shadow hashes.""" + prompt = list(range(4, 37)) + params = greedy_params() + async with routed_model( + monkeypatch, + enable_prefix_caching=True, + prefix_caching_eviction_policy=PrefixCachingEvictionPolicy.LRU, + coordinator_options={"prefix_caching_coordinator_policy": policy}, + ) as h: + direct = await h.direct(prompt, params) + await h.start() + if h.rank == 0: + for _ in range(2): + final = await asyncio.wait_for( + h.clients[0].add_request(prompt, copy.deepcopy(params)), timeout=60 + ) + assert final["generated_tokens"] == direct["generated_tokens"] + assert final["num_cached_tokens"] >= 16 + submits = [e for e in h.service.events if e["header"] == Headers.SUBMIT_REQUEST] + owners = [next(iter(e["after"].values())) for e in submits] + assert len(owners) == 2 and owners[0] == owners[1] + await h.barrier() + reused = any(step["cached"].get(1, 0) >= 16 and 1 in step["ids"] for step in h.witnesses) + assert await h.sync.all_reduce_max(int(reused)) == 1 + h.assert_retired() diff --git a/tests/unit_tests/inference/test_coordinator_sampling_pairwise.py b/tests/unit_tests/inference/test_coordinator_sampling_pairwise.py new file mode 100644 index 00000000000..fb9b5a78cc8 --- /dev/null +++ b/tests/unit_tests/inference/test_coordinator_sampling_pairwise.py @@ -0,0 +1,205 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Sampling oracles separate deterministic forward results from shared RNG draws.""" + +import asyncio +from unittest import mock + +import flashinfer +import pytest +import torch + +from tests.unit_tests.inference.coordinator_pairwise_utils import greedy_params, routed_model + +pytestmark = [pytest.mark.internal, pytest.mark.asyncio] + + +@pytest.mark.parametrize("offset_by_dp", [False, True], ids=["shared-seed", "rank-seed"]) +@pytest.mark.parametrize("logprobs_mode", ["raw_logprobs", "processed_logprobs"]) +@pytest.mark.parametrize( + "backend,filters", + [ + pytest.param("torch", (4, 0.0), id="torch-top-k"), + pytest.param("torch", (0, 0.7), id="torch-top-p"), + pytest.param("flashinfer", (4, 0.0), id="flashinfer-top-k"), + pytest.param("flashinfer", (0, 0.7), id="flashinfer-top-p"), + pytest.param("flashinfer", (4, 0.7), id="flashinfer-joint"), + ], +) +async def test_routed_stochastic_first_step_distribution( + monkeypatch, backend, logprobs_mode, filters, offset_by_dp +): + """Same prompt/history gives the same distribution, not necessarily the same draw.""" + prompt = list(range(4, 20)) + params = greedy_params( + num_tokens_to_generate=1, + temperature=0.7, + top_k=filters[0], + top_p=filters[1], + return_log_probs=True, + skip_prompt_log_probs=True, + ) + async with routed_model( + monkeypatch, + sampling_backend=backend, + logprobs_mode=logprobs_mode, + offset_sampling_seed_by_dp_rank=offset_by_dp, + ) as h: + samples = [] + sampler = h.engine.controller._sampling + sample = sampler.sample_kernel + rng = h.engine.controller.sampling_rng + owner_rank = torch.distributed.get_rank(group=h.engine.controller.dp_group) + assert sampler._rng is rng + assert rng.initial_seed() == h.config.random_seed + (owner_rank if offset_by_dp else 0) + joint = mock.Mock(wraps=flashinfer.sampling.top_k_top_p_sampling_from_logits) + if backend == "flashinfer" and all(filters): + monkeypatch.setattr(flashinfer.sampling, "top_k_top_p_sampling_from_logits", joint) + + def observe(logits, n, context, **kwargs): + indices = kwargs.get("gather_indices") + rows = logits[:n] if indices is None else logits[indices[:n]] + raw = rows.detach().clone() + distribution = sampler.log_probs_kernel(raw.float(), context).detach().clone() + rng_before = rng.get_state() + output = sample(logits, n, context, **kwargs) + assert not torch.equal( + rng_before, rng.get_state() + ), "Stochastic sampling must consume RNG" + ids = context.request_ids[ + context.paused_request_count : context.total_request_count + ].tolist() + samples.append((ids, raw.cpu(), distribution.cpu(), output[:n].clone().cpu())) + return output + + monkeypatch.setattr(sampler, "sample_kernel", observe) + await h.direct(prompt, params) + assert samples and len(samples[0][0]) == 1 + _, direct_logits, direct_distribution, _ = samples[0] + samples.clear() + await h.start() + # Direct/reference and startup work cannot witness routed joint filtering. + joint.reset_mock() + results, owners, local_ids = await h.run_paused_batch( + [(client, prompt, params) for client in range(h.dp_size)] + ) + assert local_ids == [(client, 0) for client in range(h.dp_size)] + assert len(set(owners.values())) == h.dp_size + assert all( + len(r["generated_tokens"]) == len(r["generated_log_probs"]) == 1 for r in results + ) + assert len(samples) == 1, "One-token requests sample exactly one distribution" + ids, logits, distribution, tokens = samples[0] + assert len(ids) == 1 and ids[0] != 10001 + assert any(ids[0] in step["ids"] for step in h.witnesses) + torch.testing.assert_close(logits, direct_logits, rtol=0, atol=0) + torch.testing.assert_close(distribution, direct_distribution, rtol=0, atol=0) + result = next(r for r in results if r["request_id"] == ids[0]) + token = int(tokens[0]) + assert result["generated_tokens"] == [token] + raw_distribution = torch.log_softmax(logits.float(), dim=-1) + target = raw_distribution if logprobs_mode == "raw_logprobs" else distribution + atol = 8 * torch.finfo(target.dtype).eps + assert result["generated_log_probs"][0] == pytest.approx( + target[0, token].item(), rel=0, abs=atol + ) + assert abs(raw_distribution[0, token] - distribution[0, token]) > atol + assert torch.isfinite(distribution[0, token]) + assert 1 < torch.isfinite(distribution[0]).sum() < h.config.vocab_size + evidence = [None] * h.dp_size + torch.distributed.all_gather_object(evidence, (ids[0], rng.initial_seed(), token, True)) + assert [(owner, seed) for owner, seed, _, _ in sorted(evidence)] == [ + (rank, h.config.random_seed + (rank if offset_by_dp else 0)) + for rank in range(h.dp_size) + ] + assert joint.call_count == int(backend == "flashinfer" and all(filters)) + print(f"sampling owner evidence: {sorted(evidence)}") + h.assert_retired() + + +@pytest.mark.parametrize( + "echo,detokenize,eos", + [(False, False, False), (True, True, False), (True, False, False), (False, False, True)], +) +async def test_routed_output_opt_ins(monkeypatch, echo, detokenize, eos): + """Wire prompt/text opt-ins survive routing while total length stays exact.""" + prompt = list(range(4, 20)) + params = greedy_params( + num_tokens_to_generate=None, + num_tokens_total=len(prompt) + 5, + return_prompt_tokens=echo, + detokenize_generations=detokenize, + ) + async with routed_model(monkeypatch) as h: + if eos: + direct = await h.direct(prompt, params) + params.termination_id = direct["generated_tokens"][0] + await h.start() + if h.rank == 0: + result = await asyncio.wait_for(h.clients[0].add_request(prompt, params), timeout=60) + assert result["prompt_length"] == len(prompt) + assert ( + result["generated_tokens"] == [params.termination_id] + if eos + else len(result["generated_tokens"]) == 5 + ) + assert (result["prompt_tokens"] is not None) == echo + if detokenize: + assert result["generated_text"] == h.tokenizer.detokenize( + result["generated_tokens"] + ) + else: + assert result["generated_text"] is None + assert any(0 in step["ids"] for step in h.witnesses) + await h.barrier() + h.assert_retired() + + +@pytest.mark.parametrize("keep,skip", [(False, False), (False, True), (True, False), (True, True)]) +async def test_routed_bos_stop_after_metadata_compaction(monkeypatch, keep, skip): + params = greedy_params( + num_tokens_to_generate=4, + return_prompt_tokens=True, + return_log_probs=True, + skip_prompt_log_probs=skip, + top_n_logprobs=3, + detokenize_stop_sequence=keep, + ) + async with routed_model( + monkeypatch, hidden_size=128, materialize_only_last_token_logits=False + ) as h: + h.tokenizer.bos = 3 + direct = await h.direct([3, 4, 5], params) + params.stop_words = [h.tokenizer.detokenize(direct["generated_tokens"][:2])] + params.add_BOS = True + sampler, samples = h.engine.controller._sampling, [] + sample = sampler.sample_kernel + + def observe(logits, n, context, **kwargs): + ids = context.request_ids[: context.total_request_count].tolist() + samples.append((ids, context.active_request_metadata["top_k"][: len(ids)].tolist())) + return sample(logits, n, context, **kwargs) + + monkeypatch.setattr(sampler, "sample_kernel", observe) + await h.start() + samples.clear() + short = greedy_params(num_tokens_to_generate=1, top_k=4, temperature=0.7) + (short, filler, target), owners, _ = await h.run_paused_batch( + [(0, "4 5", choice) for choice in (short, short, params)] + ) + first, other, last = (result["request_id"] for result in (short, filler, target)) + assert owners[first] == owners[last] != owners[other] + assert target["status"] == "COMPLETED" and target["prompt_tokens"][1] == [3, 4, 5] + assert target["generated_tokens"] == direct["generated_tokens"][:2] * keep + for key in ("generated_log_probs", "generated_top_n_logprobs"): + assert len(target[key]) == 2 * keep + assert target["generated_log_probs"] == [ + scores[str(token)] + for token, scores in zip(target["generated_tokens"], target["generated_top_n_logprobs"]) + ] + for key in ("prompt_log_probs", "prompt_top_n_logprobs"): + assert target[key] == direct[key] + assert await h.sync.all_reduce_max( + ([first, last], [4, 1]) in samples, ([last], [1]) in samples + ) == (1, 1) + h.assert_retired() diff --git a/tests/unit_tests/inference/test_coordinator_schedule_pairwise.py b/tests/unit_tests/inference/test_coordinator_schedule_pairwise.py new file mode 100644 index 00000000000..e1f80f57da8 --- /dev/null +++ b/tests/unit_tests/inference/test_coordinator_schedule_pairwise.py @@ -0,0 +1,120 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import asyncio +import copy + +import pytest +import torch + +from megatron.core.inference.config import PrefixCachingEvictionPolicy +from tests.unit_tests.inference.coordinator_pairwise_utils import greedy_params, routed_model, until + + +async def _submit_paused(h, prompt, params): + before = set(h.service.coordinator.request_id_to_rank) + future = h.clients[0].add_request(prompt, copy.deepcopy(params)) + await until(lambda: len(h.service.coordinator.request_id_to_rank) == len(before) + 1) + request_id = (set(h.service.coordinator.request_id_to_rank) - before).pop() + return request_id, h.service.coordinator.request_id_to_rank[request_id], future + + +@pytest.mark.internal +@pytest.mark.asyncio +async def test_routed_chunk_partition_beside_decode(monkeypatch): + prompt, params = list(range(4, 69)), greedy_params(num_tokens_to_generate=16) + async with routed_model(monkeypatch, enable_chunked_prefill=True, context_max_tokens=32) as h: + assert h.dp_size == 2 + direct = await h.direct(prompt, params) + await h.start() + parts = [] + target = [None, None] + controller = h.engine.controller + forward = controller._dynamic_step_forward_logits + + def observe(*args, **kwargs): + c = h.engine.context + rows = slice(c.paused_request_count, c.total_request_count) + tensors = (c.request_ids, c.request_query_lengths, c.request_in_prefill_status_tensor) + snapshot = [ + (count, c.num_decode_requests) + for rid, count, prefill in zip(*(tensor[rows].tolist() for tensor in tensors)) + if prefill and rid == target[0] and rid in h.engine.requests + ] + result = forward(*args, **kwargs) + parts.extend(snapshot) + return result + + monkeypatch.setattr(controller, "_dynamic_step_forward_logits", observe) + await h.pause() + if h.rank == 0: + _, decode_owner, a = await _submit_paused(h, [4, 5, 6, 7], params) + _, _, b = await _submit_paused(h, [8, 9, 10, 11], params) + target[0], target[1], future = await _submit_paused(h, prompt, params) + assert decode_owner == target[1] + torch.distributed.broadcast_object_list(target, src=0) + await h.unpause() + if h.rank == 0: + _, _, final = await asyncio.wait_for(asyncio.gather(a, b, future), timeout=60) + assert final["generated_tokens"] == direct["generated_tokens"] + await h.barrier() + owns_target = target[1] == f"mp-coord-{h.rank}".encode() + assert bool(parts) == owns_target + if parts: + assert len(parts) >= 3 and sum(count for count, _ in parts) == len(prompt) + assert all(0 < count <= 32 for count, _ in parts) + assert any(decode > 0 for _, decode in parts) + h.assert_retired() + + +@pytest.mark.internal +@pytest.mark.asyncio +@pytest.mark.parametrize("choice", ["affinity", "load", "expired"]) +async def test_routed_affinity_competes_with_load_and_expiry(monkeypatch, choice): + prompt, params = list(range(4, 37)), greedy_params() + options = dict( + prefix_caching_routing_alpha=100 if choice == "load" else 0, + prefix_cache_ttl_seconds=0 if choice == "expired" else 300, + ) + async with routed_model( + monkeypatch, + enable_prefix_caching=True, + prefix_caching_eviction_policy=PrefixCachingEvictionPolicy.LRU, + coordinator_options=options, + ) as h: + assert h.dp_size == 2 + direct = await h.direct(prompt, params) + await h.start() + await h.pause() + if h.rank == 0: + _, _, filler = await _submit_paused(h, [91, 92, 93, 94], params) + _, warm_owner, warm = await _submit_paused(h, prompt, params) + await h.unpause() + if h.rank == 0: + await asyncio.wait_for(asyncio.gather(filler, warm), timeout=60) + if choice == "expired": + await asyncio.wait_for(h.clients[0].add_request(list(range(60, 77)), params), 60) + await h.barrier() + await h.pause() + if h.rank == 0: + coordinator = h.service.coordinator + pending = [] + if choice != "expired": + _, _, blocker = await _submit_paused(h, prompt, params) + pending.append(blocker) + assert coordinator._pending_counts.sum() == 1 + assert coordinator.get_least_loaded_data_parallel_rank() != warm_owner + target, owner, future = await _submit_paused(h, prompt, params) + assert (owner == warm_owner) == (choice == "affinity") + target = [target if h.rank == 0 else None, owner if h.rank == 0 else None] + torch.distributed.broadcast_object_list(target, src=0) + await h.unpause() + if h.rank == 0: + outputs = await asyncio.wait_for(asyncio.gather(*pending, future), timeout=60) + assert outputs[-1]["generated_tokens"] == direct["generated_tokens"] + await h.barrier() + steps = [s for s in h.witnesses if target[0] in s["ids"]] + owns_target = target[1] == f"mp-coord-{h.rank}".encode() + assert bool(steps) == owns_target + if steps: + assert (max(s["cached"][target[0]] for s in steps) >= 16) == (choice == "affinity") + h.assert_retired() From f48a8832ab5053a407bd7b44d44cc80f86a2f656 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 5 Oct 2026 13:53:40 -0400 Subject: [PATCH 2/9] Keep DP adaptation within campaign test line budget Signed-off-by: Lawrence McAfee --- .../inference/test_coordinator_media_pairwise.py | 4 +--- .../inference/test_coordinator_sampling_pairwise.py | 9 +++------ 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/tests/unit_tests/inference/test_coordinator_media_pairwise.py b/tests/unit_tests/inference/test_coordinator_media_pairwise.py index ed6fed2a6d0..31ac718072c 100644 --- a/tests/unit_tests/inference/test_coordinator_media_pairwise.py +++ b/tests/unit_tests/inference/test_coordinator_media_pairwise.py @@ -266,9 +266,7 @@ def raw_media(values): ) assert repeat_route[1] == expected_repeat_owner assert different_route[1] == text_route[1] - route_info = [ - (rid, coordinator.identity_to_rank_index[owner]) for rid, owner in routes - ] + route_info = [(rid, coordinator.identity_to_rank_index[owner]) for rid, owner in routes] await h.barrier() route_payload = [route_info if h.rank == 0 else None] torch.distributed.broadcast_object_list(route_payload, src=0) diff --git a/tests/unit_tests/inference/test_coordinator_sampling_pairwise.py b/tests/unit_tests/inference/test_coordinator_sampling_pairwise.py index 09e87f23277..ae953c65f61 100644 --- a/tests/unit_tests/inference/test_coordinator_sampling_pairwise.py +++ b/tests/unit_tests/inference/test_coordinator_sampling_pairwise.py @@ -190,17 +190,14 @@ def observe(logits, n, context, **kwargs): ) first, other, last = (result["request_id"] for result in (short, filler, target)) assert owners[first] == owners[last] != owners[other] - assert target["status"] == "COMPLETED" and deserialize_tensor( - target["prompt_tokens"][1] - ).tolist() == [3, 4, 5] + assert target["status"] == "COMPLETED" + assert deserialize_tensor(target["prompt_tokens"][1]).tolist() == [3, 4, 5] assert target["generated_tokens"] == direct["generated_tokens"][:2] * keep for key in ("generated_log_probs", "generated_top_n_logprobs"): assert len(target[key]) == 2 * keep assert target["generated_log_probs"] == [ scores[str(token)] - for token, scores in zip( - target["generated_tokens"], target["generated_top_n_logprobs"] - ) + for token, scores in zip(target["generated_tokens"], target["generated_top_n_logprobs"]) ] for key in ("prompt_log_probs", "prompt_top_n_logprobs"): assert target[key] == direct[key] From 8ca7fb06dfcb0decc46863e841939ff9a4b2f5ee Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 5 Oct 2026 14:16:20 -0400 Subject: [PATCH 3/9] Preserve routed admission text oracle after flattened results Signed-off-by: Lawrence McAfee --- .../inference/test_coordinator_admission_pairwise.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/inference/test_coordinator_admission_pairwise.py b/tests/unit_tests/inference/test_coordinator_admission_pairwise.py index 33fec58e551..14f71dd907a 100644 --- a/tests/unit_tests/inference/test_coordinator_admission_pairwise.py +++ b/tests/unit_tests/inference/test_coordinator_admission_pairwise.py @@ -10,11 +10,10 @@ async def test_rejected_admission_retires_before_healthy(monkeypatch): - prompt, params = list(range(4, 20)), greedy_params() + prompt, params, bad = list(range(4, 20)), greedy_params(), 0 async with routed_model(monkeypatch, max_sequence_length=256) as h: direct = await h.direct(prompt, params) await h.start() - bad = 0 if h.rank == 0: client, events = h.clients[0], h.service.events failed = await wait_for(client.add_request([4] * 257, params), 60) @@ -33,8 +32,9 @@ async def test_rejected_admission_retires_before_healthy(monkeypatch): h.assert_retired() if h.rank == 0: final = await wait_for(client.add_request(prompt, params), 60) - for key in ("status", "generated_tokens", "generated_text"): + for key in ("status", "generated_tokens"): bad |= final[key] != direct[key] + bad |= final["generated_text"] != h.tokenizer.detokenize(direct["generated_tokens"]) submits = [e["after"] for e in events if e["header"] == Headers.SUBMIT_REQUEST] bad |= submits[-1] != {final["request_id"]: owner} replies = [ From 01991d6ef531dbc8ddabafbf1426bd267a6fcbdc Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 5 Oct 2026 15:16:56 -0400 Subject: [PATCH 4/9] test: use current inference rounders in precision coverage Signed-off-by: Lawrence McAfee --- .../unit_tests/inference/test_coordinator_backends_pairwise.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/inference/test_coordinator_backends_pairwise.py b/tests/unit_tests/inference/test_coordinator_backends_pairwise.py index a98a168b731..ec0aff34063 100644 --- a/tests/unit_tests/inference/test_coordinator_backends_pairwise.py +++ b/tests/unit_tests/inference/test_coordinator_backends_pairwise.py @@ -87,7 +87,7 @@ async def test_routed_native_model_modes_match_direct(monkeypatch, mode): else dict(fp4="e2m1", fp4_recipe="nvfp4") ) # Both direct and routed GEMMs consume the same 128-row physical batch. - for name in ("ROUNDER", "TOKEN_ROUNDER", "REQUEST_ROUNDER"): + for name in ("TOKEN_ROUNDER", "REQUEST_ROUNDER"): monkeypatch.setattr(engines.DynamicInferenceContext, name, 128) monkeypatch.setattr(engines, "set_rounder", lambda _: None) monkeypatch.setattr(engines, "TransformerConfig", lambda **kw: constructor(**(kw | overrides))) From 6ff21ce8b1640d53cb03385df4fda294a8b302ce Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 5 Oct 2026 15:29:18 -0400 Subject: [PATCH 5/9] Compare routed feature text to independently detokenized direct tokens Signed-off-by: Lawrence McAfee --- .../inference/test_coordinator_features_pairwise.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/inference/test_coordinator_features_pairwise.py b/tests/unit_tests/inference/test_coordinator_features_pairwise.py index 2d90354ce7c..1cdffc2356e 100644 --- a/tests/unit_tests/inference/test_coordinator_features_pairwise.py +++ b/tests/unit_tests/inference/test_coordinator_features_pairwise.py @@ -105,8 +105,9 @@ async def _exercise_all_owners( assert sorted(Counter(owners.values()).values()) == [requests_per_client] * harness.dp_size await harness.pause() for output in outputs: - for key in ("status", "generated_tokens", "prompt_tokens", "generated_text"): + for key in ("status", "generated_tokens", "prompt_tokens"): assert output[key] == direct[key], key + assert output["generated_text"] == harness.tokenizer.detokenize(direct["generated_tokens"]) _assert_every_owner_executed(harness, target, requests_per_client) harness.assert_retired() return target @@ -290,8 +291,7 @@ def traced_forward(*args, **kwargs): ({"fp8": True, "hidden_size": 128}, ("fp8",), ("fp8-quantized-forwards",)), ( { - # FI 0.6.6 NeoX D16 corrupts adjacent heads; exercise native D64. - "hidden_size": 256, + "hidden_size": 256, # FI 0.6.6 corrupts NeoX D16 heads; use native D64. "position_embedding_type": "rope", "use_flashinfer_fused_rope": True, }, From 79e2160ee2d5f8458dd6184e9de46470edc40c7b Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Mon, 5 Oct 2026 16:00:49 -0400 Subject: [PATCH 6/9] Observe embeddings through current vision-cache entries Signed-off-by: Lawrence McAfee --- .../unit_tests/inference/test_coordinator_media_pairwise.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/inference/test_coordinator_media_pairwise.py b/tests/unit_tests/inference/test_coordinator_media_pairwise.py index 31ac718072c..ec67d083db9 100644 --- a/tests/unit_tests/inference/test_coordinator_media_pairwise.py +++ b/tests/unit_tests/inference/test_coordinator_media_pairwise.py @@ -215,8 +215,8 @@ def raw_media(values): assert local_ids == [(0, 0), (1, 0)] assert outputs[1]["generated_tokens"] == references[0] cache_before = { - key: embedding.detach().cpu().clone() - for key, embedding in h.engine._vision_embedding_cache.items() + key: entry.embedding.detach().cpu().clone() + for key, entry in h.engine._vision_embedding_cache.items() } assert set(cache_before) == {row[2] for row in admissions} await h.pause() @@ -227,7 +227,7 @@ def raw_media(values): if allow_stale: assert set(h.engine._vision_embedding_cache) == set(cache_before) assert all( - torch.equal(h.engine._vision_embedding_cache[key].cpu(), value) + torch.equal(h.engine._vision_embedding_cache[key].embedding.cpu(), value) for key, value in cache_before.items() ) else: From 8642760c7c7330eafc70e7589099bdb56d2eff6e Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 20 Oct 2026 12:10:40 -0400 Subject: [PATCH 7/9] Update scheduler witnesses for main forward state API Signed-off-by: Lawrence McAfee --- .../unit_tests/inference/test_coordinator_lifecycle_pairwise.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/inference/test_coordinator_lifecycle_pairwise.py b/tests/unit_tests/inference/test_coordinator_lifecycle_pairwise.py index 3a313b6e855..d057dbd991a 100644 --- a/tests/unit_tests/inference/test_coordinator_lifecycle_pairwise.py +++ b/tests/unit_tests/inference/test_coordinator_lifecycle_pairwise.py @@ -104,7 +104,7 @@ def observed_replay(runner, *args, **kwargs): saver.pause.assert_called_once_with(context.TMS_TAG) else: assert not hasattr(context, "memory_buffer") - assert not h.engine.controller._async_sched_logits.is_valid + assert not h.engine.controller._async_sched_forward.is_valid await h.barrier() if h.rank == 0: assert h.service.coordinator.request_id_to_rank == route and not future.done() From 96679e0db4c7db3f49b8e3ccdd062c1b9f3b8ab9 Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 20 Oct 2026 12:40:28 -0400 Subject: [PATCH 8/9] Pin intended attention backend through current engine test configuration Signed-off-by: Lawrence McAfee --- tests/unit_tests/inference/coordinator_pairwise_utils.py | 1 + .../inference/test_coordinator_backends_pairwise.py | 7 ++++--- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/unit_tests/inference/coordinator_pairwise_utils.py b/tests/unit_tests/inference/coordinator_pairwise_utils.py index 9780c840947..456448d1b05 100644 --- a/tests/unit_tests/inference/coordinator_pairwise_utils.py +++ b/tests/unit_tests/inference/coordinator_pairwise_utils.py @@ -293,6 +293,7 @@ async def routed_model(monkeypatch, *, coordinator_options=None, engine_factory= context_block_size_tokens=16, context_buffer_size_gb=0.02, transformer_impl="transformer_engine", + flash_attention_version=3, ) settings.update(overrides) config = engine_tests.DynamicEngineTestConfig(**settings) diff --git a/tests/unit_tests/inference/test_coordinator_backends_pairwise.py b/tests/unit_tests/inference/test_coordinator_backends_pairwise.py index ec0aff34063..7223372e92c 100644 --- a/tests/unit_tests/inference/test_coordinator_backends_pairwise.py +++ b/tests/unit_tests/inference/test_coordinator_backends_pairwise.py @@ -21,8 +21,6 @@ ], ) async def test_routed_dense_backends_match_direct(monkeypatch, backend, version, symbols): - constructor = partial(engines.TransformerConfig, flash_attention_version=version or 3) - monkeypatch.setattr(engines, "TransformerConfig", constructor) build = engines.DynamicInferenceEngineTestBase._build_test_env routed = partial(features.routed_model, engine_factory=lambda config: build(config).engine) monkeypatch.setattr(features, "routed_model", routed) @@ -42,7 +40,10 @@ def observed(*args, _kernel=kernel, **kwargs): monkeypatch.setattr(attention, symbol, observed) monkeypatch.setattr(features, "_instrument_scenario_runtime", instrument_native) - options = {"transformer_impl": "inference_optimized" if version is None else "local"} + options = dict( + transformer_impl="inference_optimized" if version is None else "local", + flash_attention_version=version or 3, + ) if version == 2: options.update(context_block_size_tokens=256, context_max_tokens=256) if version == 4: From 3d8fb014ff90f3dd72751a7a054bebcf97f03b9d Mon Sep 17 00:00:00 2001 From: Lawrence McAfee Date: Tue, 20 Oct 2026 12:41:01 -0400 Subject: [PATCH 9/9] Keep backend configuration concise within the approved test budget Signed-off-by: Lawrence McAfee --- .../inference/test_coordinator_backends_pairwise.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/unit_tests/inference/test_coordinator_backends_pairwise.py b/tests/unit_tests/inference/test_coordinator_backends_pairwise.py index 7223372e92c..8936c0197b6 100644 --- a/tests/unit_tests/inference/test_coordinator_backends_pairwise.py +++ b/tests/unit_tests/inference/test_coordinator_backends_pairwise.py @@ -40,10 +40,8 @@ def observed(*args, _kernel=kernel, **kwargs): monkeypatch.setattr(attention, symbol, observed) monkeypatch.setattr(features, "_instrument_scenario_runtime", instrument_native) - options = dict( - transformer_impl="inference_optimized" if version is None else "local", - flash_attention_version=version or 3, - ) + options = dict(flash_attention_version=version or 3) + options["transformer_impl"] = "inference_optimized" if version is None else "local" if version == 2: options.update(context_block_size_tokens=256, context_max_tokens=256) if version == 4: