| |
| |
|
|
| import inspect |
| from contextlib import contextmanager |
| from dataclasses import dataclass |
| from types import SimpleNamespace |
|
|
| import pytest |
| import torch |
|
|
| from models.common.llm_runtime.lane_group import LaneGroupExecutor |
|
|
|
|
| @dataclass(frozen=True) |
| class _SamplingParams: |
| temperature: list[float] |
| top_k: torch.Tensor |
| top_p: float |
|
|
|
|
| @dataclass(frozen=True) |
| class _SingletonSamplingParams: |
| temperature: list[float] |
| top_k: torch.Tensor |
| top_p: tuple[float, ...] |
|
|
|
|
| @dataclass(frozen=True) |
| class _StatefulSamplingParams: |
| temperature: list[float] |
| top_k: torch.Tensor |
| top_p: list[float] |
| presence_penalty: list[float] |
| frequency_penalty: list[float] |
| repetition_penalty: list[float] |
| seed: list[int | None] |
| enable_log_probs: list[bool] |
| num_logprobs: list[int] |
|
|
|
|
| class _Lane: |
| def __init__(self, lane_idx, *, capacity=2): |
| self.lane_idx = lane_idx |
| self.model = SimpleNamespace(config=SimpleNamespace(max_batch_size=capacity)) |
| self.model_args = f"args-{lane_idx}" |
| self.mesh_device = f"mesh-{lane_idx}" |
| self.cache_path = f"cache-{lane_idx}" |
| self.already_warmed_up_prefill = False |
| self.calls = [] |
| self.cleanup_calls = 0 |
| self.fail_method = None |
| self.cleanup_error = None |
|
|
| def _call(self, method, kwargs): |
| self.calls.append((method, kwargs)) |
| if self.fail_method == method: |
| raise RuntimeError(f"{method} boom {self.lane_idx}") |
|
|
| def configure_paged_kv_cache(self, config): |
| self._call("configure", {"config": config}) |
| self.paged_kv_cache_config = config |
|
|
| def allocate_kv_cache(self, kv_cache_shape=None, dtype=None, num_layers=None): |
| self._call( |
| "allocate", |
| { |
| "kv_cache_shape": kv_cache_shape, |
| "dtype": dtype, |
| "num_layers": num_layers, |
| }, |
| ) |
| return f"cache-handle-{self.lane_idx}" |
|
|
| def compile_prefill( |
| self, |
| *, |
| tokens, |
| page_table, |
| prompt_lens=None, |
| start_pos=None, |
| empty_slots=None, |
| kv_cache=None, |
| sampling_params=None, |
| prompt_tokens=None, |
| output_tokens=None, |
| slot_remap=None, |
| execution=None, |
| ): |
| self._call( |
| "compile_prefill", |
| { |
| "tokens": tokens, |
| "page_table": page_table, |
| "prompt_lens": prompt_lens, |
| "start_pos": start_pos, |
| "empty_slots": empty_slots, |
| "kv_cache": kv_cache, |
| "sampling_params": sampling_params, |
| "prompt_tokens": prompt_tokens, |
| "output_tokens": output_tokens, |
| "slot_remap": slot_remap, |
| "execution": execution, |
| }, |
| ) |
|
|
| def compile_decode( |
| self, |
| *, |
| tokens, |
| start_pos, |
| page_table, |
| kv_cache=None, |
| sampling_params=None, |
| prompt_tokens=None, |
| output_tokens=None, |
| slot_remap=None, |
| reset_batch=False, |
| execution=None, |
| ): |
| self._call( |
| "compile_decode", |
| { |
| "tokens": tokens, |
| "start_pos": start_pos, |
| "page_table": page_table, |
| "kv_cache": kv_cache, |
| "sampling_params": sampling_params, |
| "prompt_tokens": prompt_tokens, |
| "output_tokens": output_tokens, |
| "slot_remap": slot_remap, |
| "reset_batch": reset_batch, |
| "execution": execution, |
| }, |
| ) |
|
|
| def warmup_model_prefill(self, *, kv_cache, can_sample_on_device, enable_trace): |
| self._call( |
| "warmup_prefill", |
| { |
| "kv_cache": kv_cache, |
| "can_sample_on_device": can_sample_on_device, |
| "enable_trace": enable_trace, |
| }, |
| ) |
| self.already_warmed_up_prefill = True |
|
|
| def warmup_model_decode( |
| self, |
| *, |
| kv_cache, |
| max_batch_size, |
| num_blocks, |
| can_sample_on_device, |
| enable_trace, |
| ): |
| self._call( |
| "warmup_decode", |
| { |
| "kv_cache": kv_cache, |
| "max_batch_size": max_batch_size, |
| "num_blocks": num_blocks, |
| "can_sample_on_device": can_sample_on_device, |
| "enable_trace": enable_trace, |
| }, |
| ) |
|
|
| def prefill_forward( |
| self, |
| tokens, |
| page_table, |
| *, |
| prompt_lens=None, |
| start_pos=None, |
| empty_slots=None, |
| kv_cache=None, |
| sampling_params=None, |
| prompt_tokens=None, |
| output_tokens=None, |
| slot_remap=None, |
| execution=None, |
| ): |
| kwargs = { |
| "tokens": tokens, |
| "page_table": page_table, |
| "prompt_lens": prompt_lens, |
| "start_pos": start_pos, |
| "empty_slots": empty_slots, |
| "kv_cache": kv_cache, |
| "sampling_params": sampling_params, |
| "prompt_tokens": prompt_tokens, |
| "output_tokens": output_tokens, |
| "slot_remap": slot_remap, |
| "execution": execution, |
| } |
| self._call("prefill", kwargs) |
| values = tokens[:, 0] + self.lane_idx * 100 |
| if sampling_params is not None: |
| return values.to(torch.int64), None |
| return values.float().view(-1, 1, 1) |
|
|
| def decode_forward( |
| self, |
| tokens, |
| start_pos, |
| page_table, |
| *, |
| kv_cache=None, |
| sampling_params=None, |
| prompt_tokens=None, |
| output_tokens=None, |
| slot_remap=None, |
| reset_batch=False, |
| read_from_device=True, |
| execution=None, |
| ): |
| kwargs = { |
| "tokens": tokens, |
| "start_pos": start_pos, |
| "page_table": page_table, |
| "kv_cache": kv_cache, |
| "sampling_params": sampling_params, |
| "prompt_tokens": prompt_tokens, |
| "output_tokens": output_tokens, |
| "slot_remap": slot_remap, |
| "reset_batch": reset_batch, |
| "read_from_device": read_from_device, |
| "execution": execution, |
| } |
| self._call("decode", kwargs) |
| values = tokens + self.lane_idx * 100 |
| if not read_from_device: |
| return f"raw-{self.lane_idx}", None |
| if sampling_params is not None: |
| return values.to(torch.int64), None |
| return values.float().view(-1, 1, 1), None |
|
|
| def can_trace_prefill(self, *, tokens, prompt_lens=None, start_pos=None, empty_slots=None): |
| self._call( |
| "can_trace_prefill", |
| { |
| "tokens": tokens, |
| "prompt_lens": prompt_lens, |
| "start_pos": start_pos, |
| "empty_slots": empty_slots, |
| }, |
| ) |
| return False |
|
|
| def read_decode_output(self, tt_out, *, async_read=False): |
| self._call("read", {"tt_out": tt_out, "async_read": async_read}) |
| host = (torch.tensor([self.lane_idx * 2, self.lane_idx * 2 + 1]), None) |
| if async_read: |
| return host, [f"event-{self.lane_idx}"] |
| return host |
|
|
| def process_decode_output_host(self, tt_out, *, is_tokens=False): |
| self._call("process", {"tt_out": tt_out, "is_tokens": is_tokens}) |
| if is_tokens: |
| return torch.tensor([self.lane_idx * 2, self.lane_idx * 2 + 1], dtype=torch.int64), None |
| return torch.full((2, 1, 3), float(self.lane_idx)), None |
|
|
| def cleanup(self): |
| self.cleanup_calls += 1 |
| if self.cleanup_error is not None: |
| raise self.cleanup_error |
|
|
|
|
| def _sampling(): |
| return _SamplingParams( |
| temperature=[0.1, 0.2, 0.3, 0.4], |
| top_k=torch.tensor([1, 2, 3, 4]), |
| top_p=0.9, |
| ) |
|
|
|
|
| _POSITIONAL = inspect.Parameter.POSITIONAL_OR_KEYWORD |
| _KEYWORD_ONLY = inspect.Parameter.KEYWORD_ONLY |
| _REQUIRED = inspect.Parameter.empty |
|
|
| _PUBLIC_SIGNATURES = { |
| "allocate_kv_cache": ( |
| ("kv_cache_shape", _POSITIONAL, None), |
| ("dtype", _POSITIONAL, None), |
| ("num_layers", _POSITIONAL, None), |
| ), |
| "compile_prefill": ( |
| ("tokens", _POSITIONAL, _REQUIRED), |
| ("page_table", _POSITIONAL, _REQUIRED), |
| ("prompt_lens", _KEYWORD_ONLY, None), |
| ("start_pos", _KEYWORD_ONLY, None), |
| ("empty_slots", _KEYWORD_ONLY, None), |
| ("kv_cache", _KEYWORD_ONLY, None), |
| ("sampling_params", _KEYWORD_ONLY, None), |
| ("prompt_tokens", _KEYWORD_ONLY, None), |
| ("output_tokens", _KEYWORD_ONLY, None), |
| ("slot_remap", _KEYWORD_ONLY, None), |
| ("execution", _KEYWORD_ONLY, None), |
| ), |
| "compile_decode": ( |
| ("tokens", _POSITIONAL, _REQUIRED), |
| ("start_pos", _POSITIONAL, _REQUIRED), |
| ("page_table", _POSITIONAL, _REQUIRED), |
| ("kv_cache", _KEYWORD_ONLY, None), |
| ("sampling_params", _KEYWORD_ONLY, None), |
| ("prompt_tokens", _KEYWORD_ONLY, None), |
| ("output_tokens", _KEYWORD_ONLY, None), |
| ("slot_remap", _KEYWORD_ONLY, None), |
| ("reset_batch", _KEYWORD_ONLY, False), |
| ("execution", _KEYWORD_ONLY, None), |
| ), |
| "warmup_model_prefill": ( |
| ("kv_cache", _KEYWORD_ONLY, _REQUIRED), |
| ("can_sample_on_device", _KEYWORD_ONLY, _REQUIRED), |
| ("enable_trace", _KEYWORD_ONLY, _REQUIRED), |
| ), |
| "warmup_model_decode": ( |
| ("kv_cache", _KEYWORD_ONLY, _REQUIRED), |
| ("max_batch_size", _KEYWORD_ONLY, _REQUIRED), |
| ("num_blocks", _KEYWORD_ONLY, _REQUIRED), |
| ("can_sample_on_device", _KEYWORD_ONLY, _REQUIRED), |
| ("enable_trace", _KEYWORD_ONLY, _REQUIRED), |
| ), |
| "prefill_forward": ( |
| ("tokens", _POSITIONAL, _REQUIRED), |
| ("page_table", _POSITIONAL, _REQUIRED), |
| ("prompt_lens", _KEYWORD_ONLY, None), |
| ("start_pos", _KEYWORD_ONLY, None), |
| ("empty_slots", _KEYWORD_ONLY, None), |
| ("kv_cache", _KEYWORD_ONLY, None), |
| ("sampling_params", _KEYWORD_ONLY, None), |
| ("prompt_tokens", _KEYWORD_ONLY, None), |
| ("output_tokens", _KEYWORD_ONLY, None), |
| ("slot_remap", _KEYWORD_ONLY, None), |
| ("execution", _KEYWORD_ONLY, None), |
| ), |
| "can_trace_prefill": ( |
| ("tokens", _KEYWORD_ONLY, _REQUIRED), |
| ("prompt_lens", _KEYWORD_ONLY, None), |
| ("start_pos", _KEYWORD_ONLY, None), |
| ("empty_slots", _KEYWORD_ONLY, None), |
| ), |
| "decode_forward": ( |
| ("tokens", _POSITIONAL, _REQUIRED), |
| ("start_pos", _POSITIONAL, _REQUIRED), |
| ("page_table", _POSITIONAL, _REQUIRED), |
| ("kv_cache", _KEYWORD_ONLY, None), |
| ("sampling_params", _KEYWORD_ONLY, None), |
| ("prompt_tokens", _KEYWORD_ONLY, None), |
| ("output_tokens", _KEYWORD_ONLY, None), |
| ("slot_remap", _KEYWORD_ONLY, None), |
| ("reset_batch", _KEYWORD_ONLY, False), |
| ("read_from_device", _KEYWORD_ONLY, True), |
| ("execution", _KEYWORD_ONLY, None), |
| ), |
| "read_decode_output": ( |
| ("tt_out", _POSITIONAL, _REQUIRED), |
| ("async_read", _KEYWORD_ONLY, False), |
| ), |
| "process_decode_output_host": ( |
| ("tt_out", _POSITIONAL, _REQUIRED), |
| ("is_tokens", _KEYWORD_ONLY, False), |
| ), |
| } |
|
|
|
|
| @pytest.mark.parametrize(("method_name", "expected"), _PUBLIC_SIGNATURES.items()) |
| def test_public_api_signatures_are_exact(method_name, expected): |
| signature = inspect.signature(getattr(LaneGroupExecutor, method_name)) |
| parameters = tuple(signature.parameters.values())[1:] |
|
|
| assert tuple((parameter.name, parameter.kind, parameter.default) for parameter in parameters) == expected |
| assert all(parameter.annotation is not inspect.Parameter.empty for parameter in parameters) |
| assert signature.return_annotation is not inspect.Signature.empty |
|
|
|
|
| def test_prefill_routes_global_slots_and_restores_source_row_order(): |
| lanes = [_Lane(0), _Lane(1)] |
| group = LaneGroupExecutor(lanes) |
| tokens = torch.tensor([[10], [11], [12]]) |
| page_table = torch.tensor([[0], [1], [2]], dtype=torch.int32) |
|
|
| output = group.prefill_forward( |
| tokens, |
| page_table, |
| empty_slots=[3, 0, 2], |
| prompt_lens=torch.tensor([1, 1, 1]), |
| sampling_params=_sampling(), |
| prompt_tokens=torch.tensor([[100], [101], [102], [103]]), |
| output_tokens=[[200], [201], [202], [203]], |
| slot_remap=torch.tensor([0, 1, 2, 3]), |
| kv_cache=["kv-0", "kv-1"], |
| ) |
|
|
| assert isinstance(output, tuple) |
| assert output[0].tolist() == [110, 11, 112] |
| assert output[0].dtype == torch.int64 |
| assert output[1] is None |
|
|
| lane0_kwargs = next(kwargs for method, kwargs in lanes[0].calls if method == "prefill") |
| lane1_kwargs = next(kwargs for method, kwargs in lanes[1].calls if method == "prefill") |
| assert lane0_kwargs["empty_slots"] == [0] |
| assert lane1_kwargs["empty_slots"] == [1, 0] |
| assert lane0_kwargs["tokens"].flatten().tolist() == [11] |
| assert lane1_kwargs["tokens"].flatten().tolist() == [10, 12] |
| assert lane0_kwargs["sampling_params"].temperature == [0.2] |
| assert lane1_kwargs["sampling_params"].temperature == [0.1, 0.3] |
| assert lane0_kwargs["sampling_params"].top_k.tolist() == [2] |
| assert lane1_kwargs["sampling_params"].top_k.tolist() == [1, 3] |
| assert lane0_kwargs["sampling_params"].top_p == 0.9 |
| assert lane0_kwargs["prompt_tokens"].flatten().tolist() == [100, -1] |
| assert lane1_kwargs["prompt_tokens"].flatten().tolist() == [102, 103] |
| assert lane0_kwargs["output_tokens"] == [[200], [-1]] |
| assert lane1_kwargs["output_tokens"] == [[202], [203]] |
| assert lane0_kwargs["slot_remap"].tolist() == [0, 1] |
| assert lane1_kwargs["slot_remap"].tolist() == [0, 1] |
| assert lane0_kwargs["kv_cache"] == "kv-0" |
| assert lane1_kwargs["kv_cache"] == "kv-1" |
|
|
|
|
| def test_prefill_log_probs_restore_source_rows_and_fill_nonrequesting_lane(): |
| lanes = [_Lane(0), _Lane(1)] |
|
|
| def prefill_with_log_probs( |
| tokens, |
| page_table, |
| *, |
| prompt_lens=None, |
| start_pos=None, |
| empty_slots=None, |
| kv_cache=None, |
| sampling_params=None, |
| execution=None, |
| ): |
| assert tokens.shape[0] == 2 |
| return torch.zeros(2, dtype=torch.int64), torch.tensor([0.25, 0.75]) |
|
|
| lanes[1].prefill_forward = prefill_with_log_probs |
| group = LaneGroupExecutor(lanes) |
|
|
| output, log_probs = group.prefill_forward( |
| tokens=torch.tensor([[10], [11], [12]]), |
| page_table=torch.tensor([[0], [1], [2]], dtype=torch.int32), |
| empty_slots=[3, 0, 2], |
| sampling_params=_sampling(), |
| ) |
|
|
| assert output.tolist() == [0, 11, 0] |
| assert isinstance(log_probs, torch.Tensor) |
| assert log_probs.tolist() == [0.25, 1.0, 0.75] |
|
|
|
|
| def test_decode_slices_fixed_lane_capacity_and_normalizes_sampled_tokens(): |
| lanes = [_Lane(0), _Lane(1)] |
| group = LaneGroupExecutor(lanes) |
|
|
| output = group.decode_forward( |
| torch.tensor([10, 11, 12, 13]), |
| torch.tensor([1, 2, 3, 4]), |
| torch.arange(4, dtype=torch.int32).view(4, 1), |
| sampling_params=_sampling(), |
| prompt_tokens=torch.tensor([[100], [101], [102], [103]]), |
| output_tokens=[[200], [201], [202], [203]], |
| slot_remap=torch.tensor([1, 0, 3, 2]), |
| kv_cache=["kv-0", "kv-1"], |
| reset_batch=True, |
| ) |
|
|
| assert output[0].tolist() == [10, 11, 112, 113] |
| assert output[0].dtype == torch.int64 |
| assert output[1] is None |
| lane0_kwargs = next(kwargs for method, kwargs in lanes[0].calls if method == "decode") |
| lane1_kwargs = next(kwargs for method, kwargs in lanes[1].calls if method == "decode") |
| assert lane0_kwargs["tokens"].tolist() == [10, 11] |
| assert lane1_kwargs["tokens"].tolist() == [12, 13] |
| assert lane0_kwargs["sampling_params"].temperature == [0.1, 0.2] |
| assert lane1_kwargs["sampling_params"].temperature == [0.3, 0.4] |
| assert lane0_kwargs["prompt_tokens"].flatten().tolist() == [100, 101] |
| assert lane1_kwargs["prompt_tokens"].flatten().tolist() == [102, 103] |
| assert lane0_kwargs["output_tokens"] == [[200], [201]] |
| assert lane1_kwargs["output_tokens"] == [[202], [203]] |
| assert lane0_kwargs["slot_remap"].tolist() == [1, 0] |
| assert lane1_kwargs["slot_remap"].tolist() == [1, 0] |
| assert lane0_kwargs["reset_batch"] is lane1_kwargs["reset_batch"] is True |
|
|
|
|
| @pytest.mark.parametrize( |
| ( |
| "layout", |
| "prefill_tokens", |
| "empty_slots", |
| "prefill_sampling", |
| "decode_tokens", |
| "decode_positions", |
| "decode_sampling", |
| "slot_remap", |
| "expected_prefill_tokens", |
| "expected_decode_tokens", |
| "expected_lane_slots", |
| "expected_lane_seeds", |
| ), |
| [ |
| pytest.param( |
| "front_packed_gathered", |
| torch.tensor([[10], [11], [12], [13]]), |
| [0, 1, 2, 3], |
| _StatefulSamplingParams( |
| temperature=[0.0, 0.8, 1.0, 1.2], |
| top_k=torch.tensor([1, 5, 7, 9]), |
| top_p=[0.0, 0.7, 0.8, 0.9], |
| presence_penalty=[0.0, 0.1, 0.2, 0.3], |
| frequency_penalty=[0.0, 0.4, 0.5, 0.6], |
| repetition_penalty=[1.0, 1.1, 1.2, 1.3], |
| seed=[100, 101, 102, 103], |
| enable_log_probs=[False, True, False, True], |
| num_logprobs=[0, 0, 0, 0], |
| ), |
| torch.tensor([10, 11, 12, 13]), |
| torch.tensor([1, 2, 3, 4]), |
| _StatefulSamplingParams( |
| temperature=[0.0, 0.8, 1.0, 1.2], |
| top_k=torch.tensor([1, 5, 7, 9]), |
| top_p=[0.0, 0.7, 0.8, 0.9], |
| presence_penalty=[0.0, 0.1, 0.2, 0.3], |
| frequency_penalty=[0.0, 0.4, 0.5, 0.6], |
| repetition_penalty=[1.0, 1.1, 1.2, 1.3], |
| seed=[100, 101, 102, 103], |
| enable_log_probs=[False, True, False, True], |
| num_logprobs=[0, 0, 0, 0], |
| ), |
| torch.tensor([1, 0, 3, 2]), |
| [10, 11, 112, 113], |
| [10, 11, 112, 113], |
| ([0, 1], [0, 1]), |
| ([100, 101], [102, 103]), |
| id="front_packed_gathered", |
| ), |
| pytest.param( |
| "stable_gap_lane", |
| torch.tensor([[11], [13]]), |
| [1, 3], |
| _StatefulSamplingParams( |
| temperature=[0.8, 1.2], |
| top_k=torch.tensor([5, 9]), |
| top_p=[0.7, 0.9], |
| presence_penalty=[0.1, 0.3], |
| frequency_penalty=[0.4, 0.6], |
| repetition_penalty=[1.1, 1.3], |
| seed=[101, 103], |
| enable_log_probs=[True, True], |
| num_logprobs=[0, 0], |
| ), |
| torch.tensor([0, 11, 0, 13]), |
| torch.tensor([-1, 2, -1, 4]), |
| _StatefulSamplingParams( |
| temperature=[0.0, 0.8, 0.0, 1.2], |
| top_k=torch.tensor([1, 5, 1, 9]), |
| top_p=[0.0, 0.7, 0.0, 0.9], |
| presence_penalty=[0.0, 0.1, 0.0, 0.3], |
| frequency_penalty=[0.0, 0.4, 0.0, 0.6], |
| repetition_penalty=[1.0, 1.1, 1.0, 1.3], |
| seed=[None, 101, None, 103], |
| enable_log_probs=[False, True, False, True], |
| num_logprobs=[0, 0, 0, 0], |
| ), |
| torch.tensor([0, 1, 2, 3]), |
| [11, 113], |
| [0, 11, 100, 113], |
| ([1], [1]), |
| ([None, 101], [None, 103]), |
| id="stable_gap_lane", |
| ), |
| ], |
| ) |
| def test_sampling_state_boundary_contract_for_gathered_and_lane_layouts( |
| layout, |
| prefill_tokens, |
| empty_slots, |
| prefill_sampling, |
| decode_tokens, |
| decode_positions, |
| decode_sampling, |
| slot_remap, |
| expected_prefill_tokens, |
| expected_decode_tokens, |
| expected_lane_slots, |
| expected_lane_seeds, |
| ): |
| """Validate layouts already constructed at the tt-metal target boundary. |
| |
| ``front_packed_gathered`` models the dense, rank-segmented payload supplied |
| after gathered-DP packing. ``stable_gap_lane`` models the fixed stable-slot |
| grid supplied by lane-DP. This deliberately does not claim to test the vLLM |
| code that constructs or gathers either payload. |
| """ |
|
|
| lanes = [_Lane(0), _Lane(1)] |
| group = LaneGroupExecutor(lanes) |
| prompt_tokens = torch.tensor([[200, -1], [210, 211], [220, -1], [230, 231]]) |
| output_tokens = [[300, -1], [310, 311], [320, -1], [330, 331]] |
| execution = [f"trace-{layout}-0", f"trace-{layout}-1"] |
|
|
| prefill_result = group.prefill_forward( |
| prefill_tokens, |
| torch.arange(len(prefill_tokens), dtype=torch.int32).view(-1, 1), |
| empty_slots=empty_slots, |
| sampling_params=prefill_sampling, |
| prompt_tokens=prompt_tokens, |
| output_tokens=output_tokens, |
| slot_remap=slot_remap, |
| execution=execution, |
| ) |
| assert prefill_result[0].tolist() == expected_prefill_tokens |
|
|
| decode_result = group.decode_forward( |
| decode_tokens, |
| decode_positions, |
| torch.arange(4, dtype=torch.int32).view(4, 1), |
| sampling_params=decode_sampling, |
| prompt_tokens=prompt_tokens, |
| output_tokens=output_tokens, |
| slot_remap=slot_remap, |
| reset_batch=True, |
| execution=execution, |
| ) |
| assert decode_result[0].tolist() == expected_decode_tokens |
|
|
| if layout == "front_packed_gathered": |
| prefill_rows = ([0, 1], [2, 3]) |
| prefill_prompt = ( |
| prompt_tokens[0:2].tolist(), |
| prompt_tokens[2:4].tolist(), |
| ) |
| prefill_output = (output_tokens[0:2], output_tokens[2:4]) |
| lane_remap = ([1, 0], [1, 0]) |
| else: |
| prefill_rows = ([0], [1]) |
| prefill_prompt = ( |
| [[-1, -1], prompt_tokens[1].tolist()], |
| [[-1, -1], prompt_tokens[3].tolist()], |
| ) |
| prefill_output = ( |
| [[-1, -1], output_tokens[1]], |
| [[-1, -1], output_tokens[3]], |
| ) |
| lane_remap = ([0, 1], [0, 1]) |
|
|
| for lane_idx, lane in enumerate(lanes): |
| prefill = next(kwargs for method, kwargs in lane.calls if method == "prefill") |
| decode = next(kwargs for method, kwargs in lane.calls if method == "decode") |
| source_rows = prefill_rows[lane_idx] |
| prefill_params = prefill["sampling_params"] |
| decode_params = decode["sampling_params"] |
| assert prefill["empty_slots"] == expected_lane_slots[lane_idx] |
| assert prefill["execution"] == execution[lane_idx] |
| assert prefill_params.temperature == [prefill_sampling.temperature[row] for row in source_rows] |
| assert prefill_params.top_k.tolist() == [int(prefill_sampling.top_k[row]) for row in source_rows] |
| assert prefill_params.top_p == [prefill_sampling.top_p[row] for row in source_rows] |
| assert prefill_params.presence_penalty == [prefill_sampling.presence_penalty[row] for row in source_rows] |
| assert prefill_params.frequency_penalty == [prefill_sampling.frequency_penalty[row] for row in source_rows] |
| assert prefill_params.repetition_penalty == [prefill_sampling.repetition_penalty[row] for row in source_rows] |
| assert prefill_params.seed == [prefill_sampling.seed[row] for row in source_rows] |
| assert prefill_params.enable_log_probs == [prefill_sampling.enable_log_probs[row] for row in source_rows] |
| assert prefill_params.num_logprobs == [prefill_sampling.num_logprobs[row] for row in source_rows] |
| assert prefill["prompt_tokens"].tolist() == prefill_prompt[lane_idx] |
| assert prefill["output_tokens"] == prefill_output[lane_idx] |
| assert prefill["slot_remap"].tolist() == lane_remap[lane_idx] |
| assert decode["execution"] == execution[lane_idx] |
| lane_slice = slice(lane_idx * 2, lane_idx * 2 + 2) |
| assert decode_params.temperature == decode_sampling.temperature[lane_slice] |
| assert decode_params.top_k.tolist() == decode_sampling.top_k[lane_slice].tolist() |
| assert decode_params.top_p == decode_sampling.top_p[lane_slice] |
| assert decode_params.presence_penalty == decode_sampling.presence_penalty[lane_slice] |
| assert decode_params.frequency_penalty == decode_sampling.frequency_penalty[lane_slice] |
| assert decode_params.repetition_penalty == decode_sampling.repetition_penalty[lane_slice] |
| assert decode_params.seed == expected_lane_seeds[lane_idx] |
| assert decode_params.enable_log_probs == decode_sampling.enable_log_probs[lane_slice] |
| assert decode_params.num_logprobs == decode_sampling.num_logprobs[lane_slice] |
| assert decode["prompt_tokens"].tolist() == prompt_tokens[lane_idx * 2 : lane_idx * 2 + 2].tolist() |
| assert decode["output_tokens"] == output_tokens[lane_idx * 2 : lane_idx * 2 + 2] |
| assert decode["slot_remap"].tolist() == lane_remap[lane_idx] |
| assert decode["reset_batch"] is True |
|
|
| raw = group.decode_forward( |
| decode_tokens, |
| decode_positions, |
| torch.arange(4, dtype=torch.int32).view(4, 1), |
| sampling_params=decode_sampling, |
| read_from_device=False, |
| ) |
| host_outputs, events = group.read_decode_output(raw, async_read=True) |
| completed = group.process_decode_output_host(host_outputs, is_tokens=True) |
| assert events == ["event-0", "event-1"] |
| assert completed[0].tolist() == [0, 1, 2, 3] |
| assert completed[1] is None |
|
|
|
|
| def test_decode_broadcasts_singleton_sampling_values_to_later_dp_lanes(): |
| lanes = [_Lane(0, capacity=1), _Lane(1, capacity=1)] |
| group = LaneGroupExecutor(lanes) |
| sampling_params = _SingletonSamplingParams( |
| temperature=[0.7], |
| top_k=torch.tensor([8]), |
| top_p=(0.4,), |
| ) |
|
|
| output = group.decode_forward( |
| tokens=torch.tensor([10, 11]), |
| start_pos=torch.tensor([1, 2]), |
| page_table=torch.arange(2, dtype=torch.int32).view(2, 1), |
| sampling_params=sampling_params, |
| ) |
|
|
| assert output[0].tolist() == [10, 111] |
| lane0_kwargs = next(kwargs for method, kwargs in lanes[0].calls if method == "decode") |
| lane1_kwargs = next(kwargs for method, kwargs in lanes[1].calls if method == "decode") |
| for lane_kwargs in (lane0_kwargs, lane1_kwargs): |
| sliced = lane_kwargs["sampling_params"] |
| assert sliced.temperature == [0.7] |
| assert sliced.top_k.tolist() == [8] |
| assert sliced.top_p == (0.4,) |
|
|
|
|
| def test_decode_without_readback_returns_lane_local_outputs(): |
| group = LaneGroupExecutor([_Lane(0), _Lane(1)]) |
|
|
| output = group.decode_forward( |
| tokens=torch.tensor([0, 1, 2, 3]), |
| start_pos=torch.tensor([0, 0, 0, 0]), |
| page_table=torch.zeros((4, 1), dtype=torch.int32), |
| read_from_device=False, |
| ) |
|
|
| assert output == [("raw-0", None), ("raw-1", None)] |
|
|
|
|
| def test_async_read_and_host_processing_run_per_lane_and_preserve_lane_order(): |
| group = LaneGroupExecutor([_Lane(0), _Lane(1)]) |
|
|
| host_outputs, events = group.read_decode_output( |
| tt_out=[("raw-0", None), ("raw-1", None)], |
| async_read=True, |
| ) |
| processed = group.process_decode_output_host(tt_out=host_outputs, is_tokens=True) |
|
|
| assert events == ["event-0", "event-1"] |
| assert processed[0].tolist() == [0, 1, 2, 3] |
| assert processed[0].dtype == torch.int64 |
| assert processed[1] is None |
|
|
|
|
| def test_warmup_replicates_lane_local_case_and_cache_to_every_lane(): |
| lanes = [_Lane(0), _Lane(1)] |
| group = LaneGroupExecutor(lanes) |
|
|
| group.warmup_model_prefill( |
| kv_cache=["kv-0", "kv-1"], |
| can_sample_on_device=True, |
| enable_trace=False, |
| ) |
| group.warmup_model_decode( |
| kv_cache=["kv-0", "kv-1"], |
| enable_trace=False, |
| max_batch_size=4, |
| num_blocks=8, |
| can_sample_on_device=True, |
| ) |
|
|
| for lane_idx, lane in enumerate(lanes): |
| prefill = next(kwargs for method, kwargs in lane.calls if method == "warmup_prefill") |
| decode = next(kwargs for method, kwargs in lane.calls if method == "warmup_decode") |
| assert prefill["kv_cache"] == f"kv-{lane_idx}" |
| assert decode["kv_cache"] == f"kv-{lane_idx}" |
| assert decode["max_batch_size"] == 2 |
| assert group.already_warmed_up_prefill |
|
|
|
|
| def test_traced_warmup_without_coordinators_fails_before_invoking_a_lane(expect_error): |
| lanes = [_Lane(0), _Lane(1)] |
| group = LaneGroupExecutor(lanes) |
|
|
| with expect_error(RuntimeError, "Every DP lane must expose the trace activation barrier"): |
| group.warmup_model_decode( |
| kv_cache=["kv-0", "kv-1"], |
| enable_trace=True, |
| max_batch_size=4, |
| num_blocks=8, |
| can_sample_on_device=False, |
| ) |
|
|
| assert [lane.calls for lane in lanes] == [[], []] |
|
|
|
|
| class _DeferredCapture: |
| def __init__(self, lane_idx, events, *, ready=True): |
| self.lane_idx = lane_idx |
| self.events = events |
| self.ready = ready |
| self.capture_pending = False |
| self.trace_activated = False |
| self.capture_calls = 0 |
|
|
| @contextmanager |
| def defer_capture(self): |
| self.events.append(("enter", self.lane_idx)) |
| try: |
| yield self |
| finally: |
| self.capture_pending = False |
| self.events.append(("exit", self.lane_idx)) |
|
|
| def register(self): |
| self.events.append(("register", self.lane_idx)) |
| self.capture_pending = self.ready |
|
|
| def activate_pending_capture(self): |
| assert self.capture_pending |
| self.events.append(("capture", self.lane_idx)) |
| self.capture_calls += 1 |
| self.trace_activated = True |
| self.capture_pending = False |
|
|
|
|
| class _BarrierLane(_Lane): |
| def __init__(self, lane_idx, events, *, ready=True, fail=False): |
| super().__init__(lane_idx) |
| self.warmup = _DeferredCapture(lane_idx, events, ready=ready) |
| self.events = events |
| self.fail = fail |
|
|
| def warmup_model_decode(self, **kwargs): |
| self.events.append(("warmup", self.lane_idx)) |
| if kwargs["enable_trace"]: |
| self.warmup.register() |
| if self.fail: |
| raise RuntimeError(f"warmup boom {self.lane_idx}") |
| return super().warmup_model_decode(**kwargs) |
|
|
|
|
| def test_traced_warmup_registers_every_lane_before_any_capture(): |
| events = [] |
| lanes = [_BarrierLane(0, events), _BarrierLane(1, events)] |
| group = LaneGroupExecutor(lanes) |
|
|
| group.warmup_model_decode( |
| kv_cache=["kv-0", "kv-1"], |
| enable_trace=True, |
| max_batch_size=4, |
| num_blocks=8, |
| can_sample_on_device=False, |
| ) |
|
|
| assert events == [ |
| ("enter", 0), |
| ("enter", 1), |
| ("warmup", 0), |
| ("register", 0), |
| ("warmup", 1), |
| ("register", 1), |
| ("capture", 0), |
| ("capture", 1), |
| ("exit", 1), |
| ("exit", 0), |
| ] |
|
|
|
|
| @pytest.mark.parametrize( |
| ("mode", "message"), |
| [ |
| ("exception", "warmup boom 1"), |
| ("asymmetry", "mixed trace activation readiness"), |
| ], |
| ) |
| def test_traced_warmup_failure_or_asymmetry_captures_no_lane(mode, message, expect_error): |
| events = [] |
| group_lanes = [ |
| _BarrierLane(0, events, ready=True), |
| _BarrierLane(1, events, ready=mode != "asymmetry", fail=mode == "exception"), |
| ] |
| group = LaneGroupExecutor(group_lanes) |
|
|
| with expect_error(RuntimeError, message): |
| group.warmup_model_decode( |
| kv_cache=["kv-0", "kv-1"], |
| enable_trace=True, |
| max_batch_size=4, |
| num_blocks=8, |
| can_sample_on_device=False, |
| ) |
|
|
| assert [lane.warmup.capture_calls for lane in group_lanes] == [0, 0] |
| assert not any(lane.warmup.capture_pending for lane in group_lanes) |
|
|
|
|
| def test_eager_warmup_does_not_enter_trace_activation_deferral(): |
| events = [] |
| lanes = [_BarrierLane(0, events), _BarrierLane(1, events)] |
| group = LaneGroupExecutor(lanes) |
|
|
| group.warmup_model_decode( |
| kv_cache=["kv-0", "kv-1"], |
| enable_trace=False, |
| max_batch_size=4, |
| num_blocks=8, |
| can_sample_on_device=False, |
| ) |
|
|
| assert events == [ |
| ("warmup", 0), |
| ("warmup", 1), |
| ] |
| assert [lane.warmup.capture_calls for lane in lanes] == [0, 0] |
|
|
|
|
| def test_compile_methods_slice_requests_to_lane_executors(): |
| lanes = [_Lane(0), _Lane(1)] |
| group = LaneGroupExecutor(lanes) |
|
|
| group.compile_prefill( |
| torch.tensor([[10], [11], [12]]), |
| torch.tensor([[0], [1], [2]], dtype=torch.int32), |
| empty_slots=[3, 0, 2], |
| kv_cache=["kv-0", "kv-1"], |
| ) |
| group.compile_decode( |
| torch.tensor([10, 11, 12, 13]), |
| torch.tensor([1, 2, 3, 4]), |
| torch.arange(4, dtype=torch.int32).view(4, 1), |
| kv_cache=["kv-0", "kv-1"], |
| ) |
|
|
| lane0_prefill = next(kwargs for method, kwargs in lanes[0].calls if method == "compile_prefill") |
| lane1_prefill = next(kwargs for method, kwargs in lanes[1].calls if method == "compile_prefill") |
| lane0_decode = next(kwargs for method, kwargs in lanes[0].calls if method == "compile_decode") |
| lane1_decode = next(kwargs for method, kwargs in lanes[1].calls if method == "compile_decode") |
| assert lane0_prefill["tokens"].flatten().tolist() == [11] |
| assert lane1_prefill["tokens"].flatten().tolist() == [10, 12] |
| assert lane0_decode["tokens"].tolist() == [10, 11] |
| assert lane1_decode["tokens"].tolist() == [12, 13] |
|
|
|
|
| def test_concrete_execution_targets_preflight_every_traced_lane_before_any_execution(expect_error): |
| class DispatchLane(_Lane): |
| def __init__(self, lane_idx, *, traceable): |
| super().__init__(lane_idx) |
| self.traceable = traceable |
| self.eager_execution = object() |
| self.traced_prefill_execution = object() |
| self.traced_decode_execution = object() |
|
|
| def can_trace_prefill(self, *, tokens, prompt_lens=None, start_pos=None): |
| self._call( |
| "can_trace_prefill", |
| { |
| "tokens": tokens, |
| "prompt_lens": prompt_lens, |
| "start_pos": start_pos, |
| }, |
| ) |
| return self.traceable |
|
|
| def prefill_forward( |
| self, |
| tokens, |
| page_table, |
| *, |
| prompt_lens=None, |
| start_pos=None, |
| empty_slots=None, |
| kv_cache=None, |
| sampling_params=None, |
| execution=None, |
| ): |
| kwargs = { |
| "tokens": tokens, |
| "page_table": page_table, |
| "prompt_lens": prompt_lens, |
| "start_pos": start_pos, |
| "empty_slots": empty_slots, |
| "kv_cache": kv_cache, |
| "sampling_params": sampling_params, |
| "execution": execution, |
| } |
| self._call("prefill", kwargs) |
| return tokens[:, :1].float().view(-1, 1, 1) |
|
|
| lanes = [DispatchLane(0, traceable=True), DispatchLane(1, traceable=False)] |
| group = LaneGroupExecutor(lanes) |
| tokens = torch.tensor([[10], [11]]) |
| page_table = torch.tensor([[0], [1]], dtype=torch.int32) |
| prompt_lens = torch.tensor([7, 9]) |
| start_pos = torch.tensor([3, 4]) |
| empty_slots = [0, 2] |
|
|
| assert not group.can_trace_prefill( |
| tokens=tokens, |
| prompt_lens=prompt_lens, |
| start_pos=start_pos, |
| empty_slots=empty_slots, |
| ) |
| kwargs = { |
| "tokens": tokens, |
| "page_table": page_table, |
| "prompt_lens": prompt_lens, |
| "start_pos": start_pos, |
| "empty_slots": empty_slots, |
| } |
| assert group.prefill_forward(execution=group.eager_execution, **kwargs).flatten().tolist() == [10.0, 11.0] |
| for lane_idx, lane in enumerate(lanes): |
| eager_call = next(call for call in lane.calls if call[0] == "prefill") |
| assert eager_call[1]["execution"] is group.eager_execution[lane_idx] |
| lane.calls.clear() |
|
|
| with expect_error(RuntimeError, "no DP lane executed"): |
| group.prefill_forward(execution=group.traced_prefill_execution, **kwargs) |
|
|
| assert [[method for method, _ in lane.calls] for lane in lanes] == [ |
| ["can_trace_prefill"], |
| ["can_trace_prefill"], |
| ] |
| for lane_idx, lane in enumerate(lanes): |
| trace_kwargs = lane.calls[0][1] |
| assert set(trace_kwargs) == {"tokens", "prompt_lens", "start_pos"} |
| assert trace_kwargs["tokens"].flatten().tolist() == [10 + lane_idx] |
| assert trace_kwargs["prompt_lens"].tolist() == [7 + 2 * lane_idx] |
| assert trace_kwargs["start_pos"].tolist() == [3 + lane_idx] |
|
|
|
|
| def test_dp1_trace_classification_slices_only_the_exact_lane_subset(): |
| class TraceLane(_Lane): |
| def can_trace_prefill(self, *, tokens, prompt_lens=None, start_pos=None): |
| self._call( |
| "can_trace_prefill", |
| { |
| "tokens": tokens, |
| "prompt_lens": prompt_lens, |
| "start_pos": start_pos, |
| }, |
| ) |
| return True |
|
|
| lane = TraceLane(0) |
| group = LaneGroupExecutor([lane]) |
|
|
| assert group.can_trace_prefill( |
| tokens=torch.tensor([[10], [11]]), |
| prompt_lens=torch.tensor([7, 9]), |
| start_pos=torch.tensor([3, 4]), |
| empty_slots=[1, 0], |
| ) |
|
|
| trace_kwargs = next(kwargs for method, kwargs in lane.calls if method == "can_trace_prefill") |
| assert set(trace_kwargs) == {"tokens", "prompt_lens", "start_pos"} |
| assert trace_kwargs["tokens"].flatten().tolist() == [10, 11] |
| assert trace_kwargs["prompt_lens"].tolist() == [7, 9] |
| assert trace_kwargs["start_pos"].tolist() == [3, 4] |
|
|
|
|
| def test_successful_traced_prefill_preflights_all_lanes_before_first_execution(): |
| events = [] |
|
|
| class TraceLane(_Lane): |
| def __init__(self, lane_idx): |
| super().__init__(lane_idx) |
| self.traced_prefill_execution = object() |
|
|
| def can_trace_prefill(self, *, tokens, prompt_lens=None, start_pos=None): |
| events.append(("preflight", self.lane_idx)) |
| return True |
|
|
| def prefill_forward(self, *args, **kwargs): |
| events.append(("execute", self.lane_idx)) |
| return super().prefill_forward(*args, **kwargs) |
|
|
| group = LaneGroupExecutor([TraceLane(0), TraceLane(1)]) |
|
|
| output = group.prefill_forward( |
| tokens=torch.tensor([[10], [11]]), |
| page_table=torch.tensor([[0], [1]], dtype=torch.int32), |
| empty_slots=[0, 2], |
| execution=group.traced_prefill_execution, |
| ) |
|
|
| assert output.flatten().tolist() == [10.0, 111.0] |
| assert events == [("preflight", 0), ("preflight", 1), ("execute", 0), ("execute", 1)] |
|
|
|
|
| @pytest.mark.parametrize( |
| "unrelated_kwargs", |
| ( |
| {"page_table": torch.zeros((1, 1), dtype=torch.int32)}, |
| {"kv_cache": ["kv-0"]}, |
| {"sampling_params": None}, |
| {"execution": (object(),)}, |
| ), |
| ) |
| def test_trace_classification_rejects_unrelated_request_fields(unrelated_kwargs, expect_error): |
| group = LaneGroupExecutor([_Lane(0)]) |
|
|
| with expect_error(TypeError, "unexpected keyword argument"): |
| group.can_trace_prefill(tokens=torch.tensor([[10]]), **unrelated_kwargs) |
|
|
| group.cleanup() |
|
|
|
|
| def test_configure_and_allocate_fan_out_with_distinct_lane_configs(): |
| lanes = [_Lane(0), _Lane(1)] |
| group = LaneGroupExecutor(lanes) |
| config = SimpleNamespace(num_blocks=8) |
|
|
| group.configure_paged_kv_cache(config) |
| handles = group.allocate_kv_cache((8, 2, 32, 64), torch.bfloat16, 2) |
|
|
| assert handles == ["cache-handle-0", "cache-handle-1"] |
| lane_configs = [next(kwargs["config"] for method, kwargs in lane.calls if method == "configure") for lane in lanes] |
| assert lane_configs[0] is not lane_configs[1] |
| assert lane_configs[0].num_blocks == lane_configs[1].num_blocks == 8 |
| for lane in lanes: |
| allocation = next(kwargs for method, kwargs in lane.calls if method == "allocate") |
| assert allocation == { |
| "kv_cache_shape": (8, 2, 32, 64), |
| "dtype": torch.bfloat16, |
| "num_layers": 2, |
| } |
|
|
|
|
| def test_side_effect_execution_target_methods_return_none_even_if_lanes_return_values(): |
| class ReturningLane(_Lane): |
| def configure_paged_kv_cache(self, config): |
| super().configure_paged_kv_cache(config) |
| return self.lane_idx |
|
|
| def compile_prefill( |
| self, |
| *, |
| tokens, |
| page_table, |
| prompt_lens=None, |
| start_pos=None, |
| empty_slots=None, |
| kv_cache=None, |
| sampling_params=None, |
| execution=None, |
| ): |
| super().compile_prefill( |
| tokens=tokens, |
| page_table=page_table, |
| prompt_lens=prompt_lens, |
| start_pos=start_pos, |
| empty_slots=empty_slots, |
| kv_cache=kv_cache, |
| sampling_params=sampling_params, |
| execution=execution, |
| ) |
| return self.lane_idx |
|
|
| def compile_decode( |
| self, |
| *, |
| tokens, |
| start_pos, |
| page_table, |
| kv_cache=None, |
| sampling_params=None, |
| reset_batch=False, |
| execution=None, |
| ): |
| super().compile_decode( |
| tokens=tokens, |
| start_pos=start_pos, |
| page_table=page_table, |
| kv_cache=kv_cache, |
| sampling_params=sampling_params, |
| reset_batch=reset_batch, |
| execution=execution, |
| ) |
| return self.lane_idx |
|
|
| def warmup_model_prefill(self, *, kv_cache, can_sample_on_device, enable_trace): |
| super().warmup_model_prefill( |
| kv_cache=kv_cache, |
| can_sample_on_device=can_sample_on_device, |
| enable_trace=enable_trace, |
| ) |
| return self.lane_idx |
|
|
| def warmup_model_decode( |
| self, |
| *, |
| kv_cache, |
| max_batch_size, |
| num_blocks, |
| can_sample_on_device, |
| enable_trace, |
| ): |
| super().warmup_model_decode( |
| kv_cache=kv_cache, |
| max_batch_size=max_batch_size, |
| num_blocks=num_blocks, |
| can_sample_on_device=can_sample_on_device, |
| enable_trace=enable_trace, |
| ) |
| return self.lane_idx |
|
|
| group = LaneGroupExecutor([ReturningLane(0), ReturningLane(1)]) |
| config = SimpleNamespace(num_blocks=8) |
| results = ( |
| group.configure_paged_kv_cache(config), |
| group.compile_prefill( |
| tokens=torch.tensor([[10], [11]]), |
| page_table=torch.tensor([[0], [1]], dtype=torch.int32), |
| empty_slots=[0, 2], |
| ), |
| group.compile_decode( |
| tokens=torch.tensor([10, 11, 12, 13]), |
| start_pos=torch.tensor([1, 2, 3, 4]), |
| page_table=torch.arange(4, dtype=torch.int32).view(4, 1), |
| ), |
| group.warmup_model_prefill( |
| kv_cache=None, |
| can_sample_on_device=False, |
| enable_trace=False, |
| ), |
| group.warmup_model_decode( |
| kv_cache=None, |
| max_batch_size=4, |
| num_blocks=8, |
| can_sample_on_device=False, |
| enable_trace=False, |
| ), |
| ) |
|
|
| assert results == (None, None, None, None, None) |
|
|
|
|
| def test_constructor_validation_failure_cleans_every_supplied_lane(expect_error): |
| lanes = [_Lane(0, capacity=2), _Lane(1, capacity=4)] |
|
|
| with expect_error(ValueError, "same fixed capacity"): |
| LaneGroupExecutor(lanes) |
|
|
| assert [lane.cleanup_calls for lane in lanes] == [1, 1] |
|
|
|
|
| def test_operation_failure_is_primary_group_becomes_terminal_and_all_lanes_cleanup(expect_error): |
| lanes = [_Lane(0), _Lane(1)] |
| lanes[1].fail_method = "allocate" |
| lanes[0].cleanup_error = RuntimeError("cleanup boom") |
| group = LaneGroupExecutor(lanes) |
|
|
| with expect_error(RuntimeError, "allocate boom 1") as exc_info: |
| group.allocate_kv_cache() |
|
|
| assert lanes[0].cleanup_calls == 1 |
| assert lanes[1].cleanup_calls == 1 |
| assert [str(error) for error in exc_info.value.cleanup_failures] == ["cleanup boom"] |
| with expect_error(RuntimeError, "terminal"): |
| group.allocate_kv_cache() |
|
|
|
|
| def test_non_null_lane_log_probs_are_aggregated_in_global_row_order(): |
| lanes = [_Lane(0), _Lane(1)] |
|
|
| def decode_with_log_probs( |
| tokens, |
| start_pos, |
| page_table, |
| *, |
| kv_cache=None, |
| sampling_params=None, |
| reset_batch=False, |
| read_from_device=True, |
| execution=None, |
| ): |
| return (tokens + 100).to(torch.int64), torch.tensor([0.25, 0.75]) |
|
|
| lanes[1].decode_forward = decode_with_log_probs |
| group = LaneGroupExecutor(lanes) |
|
|
| output, log_probs = group.decode_forward( |
| tokens=torch.tensor([0, 1, 2, 3]), |
| start_pos=torch.tensor([0, 0, 0, 0]), |
| page_table=torch.zeros((4, 1), dtype=torch.int32), |
| sampling_params=_sampling(), |
| ) |
|
|
| assert output.tolist() == [0, 1, 102, 103] |
| assert isinstance(log_probs, torch.Tensor) |
| assert log_probs.tolist() == [1.0, 1.0, 0.25, 0.75] |
|
|
|
|
| def test_cleanup_is_idempotent_and_terminal(expect_error): |
| lanes = [_Lane(0), _Lane(1)] |
| group = LaneGroupExecutor(lanes) |
|
|
| group.cleanup() |
| group.cleanup() |
|
|
| assert [lane.cleanup_calls for lane in lanes] == [1, 1] |
| with expect_error(RuntimeError, "terminal"): |
| group.prefill_forward( |
| tokens=torch.zeros((1, 1), dtype=torch.long), |
| page_table=torch.zeros((1, 1), dtype=torch.int32), |
| ) |
|
|
|
|
| def test_cleanup_retries_only_lanes_that_failed_cleanup_and_stays_terminal(expect_error): |
| class FailOnceCleanupLane(_Lane): |
| def cleanup(self): |
| self.cleanup_calls += 1 |
| if self.cleanup_calls == 1: |
| raise RuntimeError("lane cleanup failed once") |
|
|
| failing_lane = FailOnceCleanupLane(0) |
| successful_lane = _Lane(1) |
| group = LaneGroupExecutor([failing_lane, successful_lane]) |
|
|
| with expect_error(RuntimeError, "lane cleanup failed once"): |
| group.cleanup() |
|
|
| assert group.terminal |
| assert [failing_lane.cleanup_calls, successful_lane.cleanup_calls] == [1, 1] |
| with expect_error(RuntimeError, "terminal"): |
| group.allocate_kv_cache() |
|
|
| group.cleanup() |
| group.cleanup() |
|
|
| assert group.terminal |
| assert [failing_lane.cleanup_calls, successful_lane.cleanup_calls] == [2, 1] |
|
|
|
|
| def test_cleanup_retries_failed_pool_shutdown_without_recleaning_lanes(expect_error): |
| class FailOnceShutdownPool: |
| def __init__(self, pool): |
| self.pool = pool |
| self.shutdown_calls = 0 |
|
|
| def submit(self, *args, **kwargs): |
| return self.pool.submit(*args, **kwargs) |
|
|
| def shutdown(self, *, wait): |
| self.shutdown_calls += 1 |
| if self.shutdown_calls == 1: |
| raise RuntimeError("pool shutdown failed once") |
| self.pool.shutdown(wait=wait) |
|
|
| lanes = [_Lane(0), _Lane(1)] |
| group = LaneGroupExecutor(lanes) |
| assert group._output_pool is not None |
| pool = FailOnceShutdownPool(group._output_pool) |
| group._output_pool = pool |
|
|
| with expect_error(RuntimeError, "pool shutdown failed once"): |
| group.cleanup() |
|
|
| assert group.terminal |
| assert [lane.cleanup_calls for lane in lanes] == [1, 1] |
|
|
| group.cleanup() |
| group.cleanup() |
|
|
| assert pool.shutdown_calls == 2 |
| assert [lane.cleanup_calls for lane in lanes] == [1, 1] |
|
|
|
|
| def test_explicit_mesh_device_is_preserved_by_identity(): |
| mesh_device = object() |
| group = LaneGroupExecutor([_Lane(0), _Lane(1)], mesh_device=mesh_device) |
|
|
| assert group.mesh_device is mesh_device |
|
|
| group.cleanup() |
|
|
