code/models/common/tests/llm_runtime/test_lane_group.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0

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()