# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any
from unittest.mock import MagicMock, patch

import numpy as np
import pytest
import torch

from vllm.model_executor.layers.mamba.mamba_utils import (
    get_conv_copy_spec,
    get_temporal_copy_spec,
)
from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec, MambaSpec
from vllm.v1.worker.mamba_utils import (
    MambaCopyBuffers,
    MambaSpecDecodeGPUContext,
    collect_mamba_copy_meta,
    do_mamba_copy_block,
    preprocess_mamba,
)

MambaStateCopyFunc = Callable[..., Any]

# Conv + temporal copy specs, in the order the tests' MambaSpec shapes expect.
_COPY_FUNCS: tuple[MambaStateCopyFunc, ...] = (
    get_conv_copy_spec,
    get_temporal_copy_spec,
)


def postprocess_mamba(
    scheduler_output: "SchedulerOutput",
    kv_cache_config: "KVCacheConfig",
    input_batch: Any,
    requests: dict[str, Any],
    forward_context: dict[str, Any],
    mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
    copy_bufs: "MambaCopyBuffers",
):
    """CPU reference for the align-mode postprocess.

    Used as a golden against the GPU fused kernel (``postprocess_mamba_align_gpu``).
    Mirrors what the production code did before the fused kernel replaced it;
    kept here because production no longer has a CPU implementation.
    """
    assert input_batch.mamba_state_idx_cpu is not None
    num_scheduled_tokens_dict = scheduler_output.num_scheduled_tokens
    scheduled_spec_decode_tokens_dict = scheduler_output.scheduled_spec_decode_tokens
    num_accepted_tokens_cpu = input_batch.num_accepted_tokens_cpu
    mamba_state_idx_cpu = input_batch.mamba_state_idx_cpu
    mamba_group_ids = copy_bufs.mamba_group_ids
    mamba_spec = copy_bufs.mamba_spec
    copy_bufs.offset = 0
    for i, req_id in enumerate(input_batch.req_ids):
        req_state = requests[req_id]
        num_computed_tokens = req_state.num_computed_tokens
        num_draft_tokens = len(scheduled_spec_decode_tokens_dict.get(req_id, []))
        num_scheduled_tokens = num_scheduled_tokens_dict[req_id]
        num_accepted_tokens = num_accepted_tokens_cpu[i]
        num_tokens_running_state = (
            num_computed_tokens + num_scheduled_tokens - num_draft_tokens
        )
        new_num_computed_tokens = num_tokens_running_state + num_accepted_tokens - 1
        aligned_new_computed_tokens = (
            new_num_computed_tokens // mamba_spec.block_size * mamba_spec.block_size
        )
        if aligned_new_computed_tokens >= num_tokens_running_state:
            accept_token_bias = aligned_new_computed_tokens - num_tokens_running_state
            src_block_idx = mamba_state_idx_cpu[i]
            dest_block_idx = aligned_new_computed_tokens // mamba_spec.block_size - 1
            collect_mamba_copy_meta(
                copy_bufs,
                kv_cache_config,
                mamba_state_copy_funcs,
                mamba_group_ids,
                src_block_idx,
                dest_block_idx,
                accept_token_bias,
                req_state,
                forward_context,
            )
            if src_block_idx == dest_block_idx:
                num_accepted_tokens_cpu[i] = 1
    do_mamba_copy_block(copy_bufs)


def _make_scheduler_output(
    finished_req_ids: set[str],
    preempted_req_ids: set[str] | None,
    resumed_req_ids: set[str],
) -> SchedulerOutput:
    cached = CachedRequestData.make_empty()
    cached.resumed_req_ids = resumed_req_ids
    return SchedulerOutput(
        scheduled_new_reqs=[],
        scheduled_cached_reqs=cached,
        num_scheduled_tokens={},
        total_num_scheduled_tokens=0,
        scheduled_spec_decode_tokens={},
        scheduled_encoder_inputs={},
        num_common_prefix_blocks=[],
        finished_req_ids=finished_req_ids,
        free_encoder_mm_hashes=[],
        preempted_req_ids=preempted_req_ids,
    )


def test_resumed_req_ids_cleared_from_mamba_state_idx():
    """When a request is force-preempted (e.g. reset_prefix_cache),
    it appears in resumed_req_ids but NOT in preempted_req_ids.
    preprocess_mamba must still clear its mamba_state_idx entry,
    otherwise stale indices can point beyond the new block allocation.
    """
    spec = MagicMock(block_size=64, num_speculative_blocks=0)
    cache_config = MagicMock(enable_prefix_caching=True)
    input_batch = MagicMock(req_ids=[])
    copy_bufs = MagicMock(mamba_group_ids=[0], mamba_spec=spec)

    mamba_state_idx: dict[str, int] = {
        "finished": 1,
        "preempted": 2,
        "resumed": 3,  # only in resumed_req_ids, NOT in preempted
        "keep": 99,
    }
    sched = _make_scheduler_output(
        finished_req_ids={"finished"},
        preempted_req_ids={"preempted"},
        resumed_req_ids={"resumed"},
    )

    with patch(
        "vllm.v1.worker.mamba_utils.get_mamba_groups",
        return_value=([0], spec),
    ):
        preprocess_mamba(
            sched,
            MagicMock(),  # kv_cache_config
            cache_config,
            mamba_state_idx,
            input_batch,
            {},  # requests
            {},  # forward_context
            (),  # mamba_state_copy_funcs
            copy_bufs,
        )

    assert mamba_state_idx == {"keep": 99}


# -----------------------------------------------------------------------------
# Golden tests for postprocess_mamba_fused_kernel
# -----------------------------------------------------------------------------


@dataclass
class _TestConfig:
    """Common test configuration for fused kernel tests."""

    block_size: int = 16
    num_blocks: int = 32
    num_layers: int = 2
    num_reqs: int = 4
    max_num_reqs: int = 8
    # Conv state shape: [num_blocks, conv_width, inner_dim]
    conv_width: int = 4
    conv_inner_dim: int = 64
    # Temporal state shape: [num_blocks, state_dim]
    temporal_state_dim: int = 128
    dtype: torch.dtype = torch.float16


class _MockCpuGpuBuffer:
    """Mock CpuGpuBuffer for testing without pinned memory."""

    def __init__(self, size: int, dtype: torch.dtype, device: torch.device):
        self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
        self.gpu = torch.zeros(size, dtype=dtype, device=device)
        self.np = self.cpu.numpy()

    def copy_to_gpu(self, n: int | None = None) -> torch.Tensor:
        if n is None:
            return self.gpu.copy_(self.cpu, non_blocking=True)
        return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True)


def _make_postprocess_scheduler_output(
    req_ids: list[str],
    num_scheduled_tokens: dict[str, int],
    scheduled_spec_decode_tokens: dict[str, list] | None = None,
) -> SchedulerOutput:
    """Create a minimal SchedulerOutput for postprocess testing."""
    cached = CachedRequestData.make_empty()
    return SchedulerOutput(
        scheduled_new_reqs=[],
        scheduled_cached_reqs=cached,
        num_scheduled_tokens=num_scheduled_tokens,
        total_num_scheduled_tokens=sum(num_scheduled_tokens.values()),
        scheduled_spec_decode_tokens=scheduled_spec_decode_tokens or {},
        scheduled_encoder_inputs={},
        num_common_prefix_blocks=[],
        finished_req_ids=set(),
        free_encoder_mm_hashes=[],
        preempted_req_ids=set(),
    )


def _make_mock_attention(
    conv_state: torch.Tensor, temporal_state: torch.Tensor
) -> MagicMock:
    """Create a mock attention object with kv_cache."""
    attention = MagicMock()
    attention.kv_cache = [conv_state, temporal_state]
    return attention


def _make_dual_states(
    cfg: "_TestConfig",
    layer_names: list[str],
    device: torch.device,
    *,
    num_blocks: int | None = None,
) -> tuple[
    list[torch.Tensor],
    list[torch.Tensor],
    list[torch.Tensor],
    list[torch.Tensor],
    dict[str, MagicMock],
    dict[str, MagicMock],
]:
    """Allocate conv+temporal state tensors for the Python path, clone them for
    the GPU path, and build matching ``forward_context`` dicts for both.

    Returns ``(conv_py, temporal_py, conv_gpu, temporal_gpu, fwd_py, fwd_gpu)``
    where the four state lists are parallel to ``layer_names``.
    """
    n_blocks = num_blocks if num_blocks is not None else cfg.num_blocks
    conv_py = [
        torch.randn(
            n_blocks,
            cfg.conv_width,
            cfg.conv_inner_dim,
            dtype=cfg.dtype,
            device=device,
        )
        for _ in layer_names
    ]
    temporal_py = [
        torch.randn(n_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device)
        for _ in layer_names
    ]
    conv_gpu = [s.clone() for s in conv_py]
    temporal_gpu = [s.clone() for s in temporal_py]
    fwd_py = {
        name: _make_mock_attention(c, t)
        for name, c, t in zip(layer_names, conv_py, temporal_py)
    }
    fwd_gpu = {
        name: _make_mock_attention(c, t)
        for name, c, t in zip(layer_names, conv_gpu, temporal_gpu)
    }
    return conv_py, temporal_py, conv_gpu, temporal_gpu, fwd_py, fwd_gpu


def _make_dual_layer_state(
    cfg: "_TestConfig",
    device: torch.device,
    *,
    num_blocks: int | None = None,
    layer_name: str = "layer_0",
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    dict[str, MagicMock],
    dict[str, MagicMock],
]:
    """Single-layer convenience form of ``_make_dual_states``."""
    conv_py, temporal_py, conv_gpu, temporal_gpu, fwd_py, fwd_gpu = _make_dual_states(
        cfg, [layer_name], device, num_blocks=num_blocks
    )
    return conv_py[0], temporal_py[0], conv_gpu[0], temporal_gpu[0], fwd_py, fwd_gpu


def _make_kv_cache_config(cfg: _TestConfig, layer_names: list[str]) -> KVCacheConfig:
    """Create a KVCacheConfig with mamba groups."""
    mamba_spec = MambaSpec(
        block_size=cfg.block_size,
        shapes=(
            (cfg.conv_width, cfg.conv_inner_dim),
            (cfg.temporal_state_dim,),
        ),
        dtypes=(cfg.dtype, cfg.dtype),
        mamba_cache_mode="all",
    )
    group = KVCacheGroupSpec(
        layer_names=layer_names,
        kv_cache_spec=mamba_spec,
    )
    return KVCacheConfig(
        num_blocks=cfg.num_blocks,
        kv_cache_tensors=[],
        kv_cache_groups=[group],
    )


def _make_input_batch(
    req_ids: list[str],
    num_accepted_tokens: list[int],
    mamba_state_idx: list[int],
) -> MagicMock:
    """Create a mock GPUInputBatch."""
    batch = MagicMock()
    batch.req_ids = req_ids
    batch.req_id_to_index = {rid: i for i, rid in enumerate(req_ids)}
    # Use numpy arrays so modifications persist
    batch.num_accepted_tokens_cpu = np.array(num_accepted_tokens, dtype=np.int32)
    batch.mamba_state_idx_cpu = np.array(mamba_state_idx, dtype=np.int32)
    return batch


def _make_requests(
    req_ids: list[str],
    num_computed_tokens: list[int],
    block_ids_per_req: list[list[int]],
) -> dict[str, MagicMock]:
    """Create mock CachedRequestState objects."""
    requests = {}
    for i, req_id in enumerate(req_ids):
        req = MagicMock()
        req.num_computed_tokens = num_computed_tokens[i]
        req.block_ids = {0: block_ids_per_req[i]}  # group_id=0
        requests[req_id] = req
    return requests


def _make_copy_bufs(
    cfg: _TestConfig, kv_cache_config: KVCacheConfig, device: torch.device
) -> MambaCopyBuffers:
    """Create MambaCopyBuffers for the Python path."""

    def make_buffer(n, dtype):
        return _MockCpuGpuBuffer(n, dtype, device)

    return MambaCopyBuffers.create(
        max_num_reqs=cfg.max_num_reqs,
        kv_cache_config=kv_cache_config,
        copy_funcs=(get_conv_copy_spec, get_temporal_copy_spec),
        make_buffer=make_buffer,
    )


def _make_gpu_ctx(
    cfg: _TestConfig, kv_cache_config: KVCacheConfig, device: torch.device
) -> MambaSpecDecodeGPUContext:
    """Create MambaSpecDecodeGPUContext for the GPU path."""

    def make_buffer(n, dtype):
        return _MockCpuGpuBuffer(n, dtype, device)

    return MambaSpecDecodeGPUContext.create(
        max_num_reqs=cfg.max_num_reqs,
        kv_cache_config=kv_cache_config,
        num_state_types=2,
        device=device,
        make_buffer=make_buffer,
    )


def _run_gpu_postprocess(
    gpu_ctx: MambaSpecDecodeGPUContext,
    *,
    kv_cache_config: KVCacheConfig,
    forward_context: dict[str, Any],
    copy_funcs: tuple,
    block_table: torch.Tensor,
    req_ids: list[str],
    num_accepted_tokens: list[int],
    mamba_state_idx: list[int],
    num_scheduled_tokens: dict[str, int],
    num_computed_tokens: list[int],
    num_draft_tokens: dict[str, int],
    device: torch.device,
) -> None:
    """Initialize the GPU context against `block_table`, run the fused
    postprocess kernel for `req_ids`, and synchronize."""

    def t(values):
        return torch.tensor(values, dtype=torch.int32, device=device)

    gpu_ctx.initialize_from_forward_context(
        kv_cache_config, forward_context, copy_funcs, [block_table]
    )
    gpu_ctx.run_fused_postprocess(
        num_reqs=len(req_ids),
        num_accepted_tokens_gpu=t(num_accepted_tokens),
        mamba_state_idx_gpu=t(mamba_state_idx),
        num_scheduled_tokens_gpu=t([num_scheduled_tokens[r] for r in req_ids]),
        num_computed_tokens_gpu=t(num_computed_tokens),
        num_draft_tokens_gpu=t([num_draft_tokens.get(r, 0) for r in req_ids]),
    )
    torch.accelerator.synchronize()


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
class TestPostprocessMambaFusedKernel:
    """Tests for postprocess_mamba_fused_kernel comparing GPU vs CPU paths."""

    @pytest.fixture
    def device(self):
        return torch.device("cuda:0")

    @pytest.fixture
    def test_config(self):
        return _TestConfig()

    def test_matches_python_postprocess_mamba(self, device, test_config):
        """
        Golden test: GPU kernel produces identical results to Python impl.

        This test:
        1. Sets up identical initial state for both paths
        2. Runs Python postprocess_mamba (modifies states via batch_memcpy)
        3. Runs GPU fused kernel (modifies states directly)
        4. Compares resulting state tensors and num_accepted_tokens
        """
        cfg = test_config
        torch.manual_seed(42)

        # Test scenario: 4 requests with different copy conditions
        # Copy needed when: aligned_new_computed >= num_tokens_running_state
        # where: num_tokens_running_state = num_computed + num_scheduled - num_draft
        #        new_num_computed = num_tokens_running_state + num_accepted - 1
        #        aligned_new_computed = (new_num_computed // block_size) * block_size
        req_ids = ["req_0", "req_1", "req_2", "req_3"]

        # Configure requests so some need copies, some don't
        # block_size = 16
        # req_0: running=60+5-2=63, new=63+3-1=65, aligned=64 >= 63 -> COPY
        # req_1: running=30+3-0=33, new=33+2-1=34, aligned=32 < 33 -> NO COPY
        # req_2: running=45+8-3=50, new=50+4-1=53, aligned=48 < 50 -> NO COPY
        # req_3: running=10+6-0=16, new=16+2-1=17, aligned=16 >= 16 -> COPY
        num_computed_tokens = [60, 30, 45, 10]
        num_scheduled_tokens = {"req_0": 5, "req_1": 3, "req_2": 8, "req_3": 6}
        num_draft_tokens = {"req_0": 2, "req_1": 0, "req_2": 3, "req_3": 0}
        num_accepted_tokens = [3, 2, 4, 2]
        mamba_state_idx = [3, 1, 2, 0]  # source block indices

        # Block IDs for each request (simulate block table)
        block_ids_per_req = [
            list(range(8)),  # req_0: blocks 0-7
            list(range(8, 16)),  # req_1: blocks 8-15
            list(range(16, 24)),  # req_2: blocks 16-23
            list(range(24, 32)),  # req_3: blocks 24-31
        ]

        layer_names = [f"layer_{i}" for i in range(cfg.num_layers)]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        (
            conv_states_py,
            temporal_states_py,
            conv_states_gpu,
            temporal_states_gpu,
            forward_context_py,
            forward_context_gpu,
        ) = _make_dual_states(cfg, layer_names, device)

        # --- Run Python path ---
        scheduler_output = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        input_batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)

        postprocess_mamba(
            scheduler_output,
            kv_cache_config,
            input_batch_py,
            requests,
            forward_context_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # --- Run GPU path ---
        gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)

        num_reqs = len(req_ids)
        max_blocks = max(len(b) for b in block_ids_per_req)
        block_table_gpu = torch.zeros(
            num_reqs, max_blocks, dtype=torch.int32, device=device
        )
        for i, block_ids in enumerate(block_ids_per_req):
            block_table_gpu[i, : len(block_ids)] = torch.tensor(
                block_ids, dtype=torch.int32
            )
        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # --- Compare results ---
        # 1. Compare state tensors
        for i in range(cfg.num_layers):
            torch.testing.assert_close(
                conv_states_gpu[i],
                conv_states_py[i],
                msg=f"Conv state mismatch at layer {i}",
            )
            torch.testing.assert_close(
                temporal_states_gpu[i],
                temporal_states_py[i],
                msg=f"Temporal state mismatch at layer {i}",
            )

        # 2. Compare num_accepted_tokens updates
        expected_accepted = torch.tensor(
            input_batch_py.num_accepted_tokens_cpu[:num_reqs],
            dtype=torch.int32,
            device=device,
        )
        torch.testing.assert_close(
            gpu_ctx.num_accepted_tokens_out[:num_reqs],
            expected_accepted,
            msg="num_accepted_tokens mismatch",
        )

    def test_no_copy_when_not_needed(self, device, test_config):
        """Kernel should not modify state when no copy is needed."""
        cfg = test_config
        torch.manual_seed(123)

        # Single request where no copy is needed:
        # running = 30 + 3 = 33, new = 33 + 1 - 1 = 33, aligned = 32 < 33
        req_ids = ["req_0"]
        num_computed_tokens = [30]
        num_scheduled_tokens = {"req_0": 3}
        num_draft_tokens: dict[str, int] = {}
        num_accepted_tokens = [1]
        mamba_state_idx = [1]
        block_ids_per_req = [list(range(8))]

        layer_names = ["layer_0"]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        # Create state tensor
        conv_state = torch.randn(
            cfg.num_blocks,
            cfg.conv_width,
            cfg.conv_inner_dim,
            dtype=cfg.dtype,
            device=device,
        )
        temporal_state = torch.randn(
            cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device
        )

        # Clone to verify no modification
        conv_state_orig = conv_state.clone()
        temporal_state_orig = temporal_state.clone()

        forward_context = {"layer_0": _make_mock_attention(conv_state, temporal_state)}

        gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)

        num_reqs = len(req_ids)
        block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
        block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # State should be unchanged
        torch.testing.assert_close(conv_state, conv_state_orig)
        torch.testing.assert_close(temporal_state, temporal_state_orig)

    @pytest.mark.parametrize("num_reqs", [1, 2, 8, 16])
    def test_various_batch_sizes(self, device, test_config, num_reqs):
        """Verify kernel works correctly with different batch sizes."""
        cfg = _TestConfig(max_num_reqs=max(16, num_reqs))
        torch.manual_seed(456)

        req_ids = [f"req_{i}" for i in range(num_reqs)]
        # All requests will trigger a copy
        num_computed_tokens = [60] * num_reqs
        num_scheduled_tokens = {r: 5 for r in req_ids}
        num_draft_tokens = {r: 0 for r in req_ids}
        num_accepted_tokens = [3] * num_reqs
        mamba_state_idx = [3] * num_reqs
        # Each request gets unique blocks
        block_ids_per_req = [list(range(i * 8, (i + 1) * 8)) for i in range(num_reqs)]

        # Ensure we have enough blocks
        total_blocks = num_reqs * 8
        cfg = _TestConfig(num_blocks=total_blocks, max_num_reqs=max(16, num_reqs))

        layer_names = ["layer_0"]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        (
            conv_state_py,
            temporal_state_py,
            conv_state_gpu,
            temporal_state_gpu,
            forward_context_py,
            forward_context_gpu,
        ) = _make_dual_layer_state(cfg, device)

        # Run Python path
        scheduler_output = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        input_batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)

        postprocess_mamba(
            scheduler_output,
            kv_cache_config,
            input_batch_py,
            requests,
            forward_context_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # Run GPU path
        gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)

        max_blocks_per_req = 8
        block_table_gpu = torch.zeros(
            num_reqs, max_blocks_per_req, dtype=torch.int32, device=device
        )
        for i, block_ids in enumerate(block_ids_per_req):
            block_table_gpu[i, : len(block_ids)] = torch.tensor(
                block_ids, dtype=torch.int32
            )

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # Compare results
        torch.testing.assert_close(
            conv_state_gpu, conv_state_py, msg="Conv state mismatch"
        )
        torch.testing.assert_close(
            temporal_state_gpu, temporal_state_py, msg="Temporal state mismatch"
        )

    def test_block_table_with_realistic_stride(self, device, test_config):
        """
        Test kernel with realistic block table strides.

        In real usage, the block table is pre-allocated with shape
        [max_num_reqs, max_num_blocks_per_req] and then sliced to
        [:num_reqs]. This means stride(0) = max_num_blocks_per_req,
        which is typically much larger than the actual blocks used.

        This test verifies the kernel handles non-tight strides correctly,
        catching bugs where stride is incorrectly treated as bytes vs elements.
        """
        cfg = test_config
        torch.manual_seed(789)

        # Use multiple requests to exercise stride-based indexing
        num_reqs = 4
        req_ids = [f"req_{i}" for i in range(num_reqs)]

        # All requests trigger copies (same setup as test_various_batch_sizes)
        num_computed_tokens = [60] * num_reqs
        num_scheduled_tokens = {r: 5 for r in req_ids}
        num_draft_tokens: dict[str, int] = {}
        num_accepted_tokens = [3] * num_reqs
        mamba_state_idx = [3] * num_reqs

        # Each request uses only 8 blocks, but we allocate much more
        blocks_used_per_req = 8
        block_ids_per_req = [
            list(range(i * blocks_used_per_req, (i + 1) * blocks_used_per_req))
            for i in range(num_reqs)
        ]

        total_blocks = num_reqs * blocks_used_per_req
        cfg = _TestConfig(num_blocks=total_blocks, max_num_reqs=max(16, num_reqs))

        layer_names = ["layer_0"]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        (
            conv_state_py,
            temporal_state_py,
            conv_state_gpu,
            temporal_state_gpu,
            forward_context_py,
            forward_context_gpu,
        ) = _make_dual_layer_state(cfg, device)

        # Run Python path
        scheduler_output = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        input_batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)

        postprocess_mamba(
            scheduler_output,
            kv_cache_config,
            input_batch_py,
            requests,
            forward_context_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # Run GPU path with REALISTIC block table stride
        gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)

        # KEY DIFFERENCE: Create a large block table like real code does
        # Real system has max_num_blocks_per_req >> blocks actually used
        max_num_reqs_full = 16
        max_blocks_per_req_full = 512  # Much larger than blocks_used_per_req=8

        # Allocate full-size table (simulates pre-allocated CpuGpuBuffer)
        block_table_full = torch.zeros(
            max_num_reqs_full, max_blocks_per_req_full, dtype=torch.int32, device=device
        )

        # Fill in actual block IDs (only first few columns used)
        for i, block_ids in enumerate(block_ids_per_req):
            block_table_full[i, : len(block_ids)] = torch.tensor(
                block_ids, dtype=torch.int32
            )

        # Slice like real code: block_table.gpu[:num_reqs]
        # This preserves stride(0) = 512, not 8!
        block_table_gpu = block_table_full[:num_reqs]

        # Verify stride is large (the key property we're testing)
        assert block_table_gpu.stride(0) == max_blocks_per_req_full, (
            f"Expected stride {max_blocks_per_req_full}, "
            f"got {block_table_gpu.stride(0)}"
        )

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # Compare results - this will fail if stride handling is incorrect
        torch.testing.assert_close(
            conv_state_gpu,
            conv_state_py,
            msg="Conv state mismatch - possible stride bug in kernel",
        )
        torch.testing.assert_close(
            temporal_state_gpu,
            temporal_state_py,
            msg="Temporal state mismatch - possible stride bug in kernel",
        )

    def test_src_addr_equals_dst_addr_skips_copy_and_sets_accepted_to_1(
        self, device, test_config
    ):
        """
        Test the ``src_addr == dst_addr`` early-return path in
        postprocess_mamba_fused_kernel matches Python behavior.

        When src_addr == dst_addr (source and destination memory addresses are
        identical), both implementations should:
        1. Skip the copy (state unchanged)
        2. Set num_accepted_tokens to 1

        This condition occurs when:
        - src_block_idx == dest_block_idx (same logical block)
        - accept_token_bias == 0 (no offset within the block)

        Python reference (collect_mamba_copy_meta):
            if src_block_idx == dest_block_idx and accept_token_bias == 0:
                return  # No copy added

        Python reference (postprocess_mamba):
            if src_block_idx == dest_block_idx:
                num_accepted_tokens_cpu[i] = 1

        Test setup (block_size=16):
        - num_tokens_running_state = 30 + 2 - 0 = 32
        - new_num_computed = 32 + 1 - 1 = 32
        - aligned_new_computed = 32
        - accept_token_bias = 32 - 32 = 0
        - dest_block_idx = 32 // 16 - 1 = 1
        - src_block_idx = 1 (set explicitly)
        """
        cfg = test_config
        torch.manual_seed(1001)

        req_ids = ["req_0"]
        num_computed_tokens = [30]
        num_scheduled_tokens = {"req_0": 2}
        num_draft_tokens: dict[str, int] = {}
        num_accepted_tokens = [5]  # Initial value, should become 1
        mamba_state_idx = [1]  # src_block_idx = 1 = dest_block_idx
        block_ids_per_req = [list(range(8))]

        layer_names = ["layer_0"]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        (
            conv_state_py,
            temporal_state_py,
            conv_state_gpu,
            temporal_state_gpu,
            forward_context_py,
            forward_context_gpu,
        ) = _make_dual_layer_state(cfg, device)

        # Also clone to verify no modification
        conv_state_orig = conv_state_py.clone()
        temporal_state_orig = temporal_state_py.clone()

        # --- Run Python path ---
        scheduler_output = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        input_batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)

        postprocess_mamba(
            scheduler_output,
            kv_cache_config,
            input_batch_py,
            requests,
            forward_context_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # --- Run GPU path ---
        gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)

        num_reqs = len(req_ids)
        block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
        block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # --- Verify Python behavior (ground truth) ---
        # State should be unchanged (no copy when src_addr == dst_addr)
        torch.testing.assert_close(
            conv_state_py,
            conv_state_orig,
            msg="Python: Conv state should be unchanged when src==dst",
        )
        torch.testing.assert_close(
            temporal_state_py,
            temporal_state_orig,
            msg="Python: Temporal state should be unchanged when src==dst",
        )
        # num_accepted_tokens should be 1
        assert input_batch_py.num_accepted_tokens_cpu[0] == 1, (
            f"Python: num_accepted_tokens should be 1, "
            f"got {input_batch_py.num_accepted_tokens_cpu[0]}"
        )

        # --- Verify GPU matches Python ---
        torch.testing.assert_close(
            conv_state_gpu,
            conv_state_py,
            msg="GPU conv state should match Python",
        )
        torch.testing.assert_close(
            temporal_state_gpu,
            temporal_state_py,
            msg="GPU temporal state should match Python",
        )
        expected_accepted = torch.tensor(
            input_batch_py.num_accepted_tokens_cpu[:num_reqs],
            dtype=torch.int32,
            device=device,
        )
        torch.testing.assert_close(
            gpu_ctx.num_accepted_tokens_out[:num_reqs],
            expected_accepted,
            msg="GPU num_accepted_tokens should match Python",
        )

    def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1(
        self, device, test_config
    ):
        """
        Test the ``src_block_idx == dest_block_idx`` post-copy update in
        postprocess_mamba_fused_kernel matches Python behavior.

        When src_block_idx == dest_block_idx but accept_token_bias > 0, both
        implementations should:
        1. Perform the copy (src_addr != dst_addr due to offset)
        2. Set num_accepted_tokens to 1 AFTER the copy

        Python reference (postprocess_mamba):
            if src_block_idx == dest_block_idx:
                num_accepted_tokens_cpu[i] = 1

        For conv states: copies state[block, offset:] to
            state[block, :] (shifted window)
        For temporal states: copies state[block_ids[src_idx + offset]] to
            state[block_ids[dest_idx]]

        Test setup (block_size=16):
        - num_tokens_running_state = 30 + 1 - 0 = 31
        - new_num_computed = 31 + 2 - 1 = 32
        - aligned_new_computed = 32
        - accept_token_bias = 32 - 31 = 1 (> 0, so copy happens)
        - dest_block_idx = 32 // 16 - 1 = 1
        - src_block_idx = 1 (set explicitly, == dest_block_idx)
        """
        cfg = test_config
        torch.manual_seed(1002)

        req_ids = ["req_0"]
        num_computed_tokens = [30]
        num_scheduled_tokens = {"req_0": 1}
        num_draft_tokens: dict[str, int] = {}
        num_accepted_tokens = [2]  # Results in accept_token_bias = 1
        mamba_state_idx = [1]  # src_block_idx = 1 = dest_block_idx
        block_ids_per_req = [list(range(8))]

        layer_names = ["layer_0"]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        (
            conv_state_py,
            temporal_state_py,
            conv_state_gpu,
            temporal_state_gpu,
            forward_context_py,
            forward_context_gpu,
        ) = _make_dual_layer_state(cfg, device)

        # Clone to verify modification
        conv_state_orig = conv_state_py.clone()
        temporal_state_orig = temporal_state_py.clone()

        # --- Run Python path ---
        scheduler_output = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        input_batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)

        postprocess_mamba(
            scheduler_output,
            kv_cache_config,
            input_batch_py,
            requests,
            forward_context_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # --- Run GPU path ---
        gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)

        num_reqs = len(req_ids)
        block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
        block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # --- Verify Python behavior (ground truth) ---
        dest_block_id = block_ids_per_req[0][1]  # dest_block_idx = 1

        # Conv state should be modified (shifted copy within block)
        conv_changed = not torch.allclose(
            conv_state_py[dest_block_id], conv_state_orig[dest_block_id]
        )
        assert conv_changed, (
            "Python: Conv state should be modified when accept_token_bias > 0"
        )

        # Temporal state should be modified (copy from different block)
        src_block_id_temporal = block_ids_per_req[0][2]  # actual_src_block_idx = 2
        dest_block_id_temporal = block_ids_per_req[0][1]  # dest_block_idx = 1
        torch.testing.assert_close(
            temporal_state_py[dest_block_id_temporal],
            temporal_state_orig[src_block_id_temporal],
            msg="Python: Temporal state copy should have happened",
        )

        # num_accepted_tokens should be 1
        assert input_batch_py.num_accepted_tokens_cpu[0] == 1, (
            f"Python: num_accepted_tokens should be 1, "
            f"got {input_batch_py.num_accepted_tokens_cpu[0]}"
        )

        # --- Verify GPU matches Python ---
        torch.testing.assert_close(
            conv_state_gpu,
            conv_state_py,
            msg="GPU conv state should match Python",
        )
        torch.testing.assert_close(
            temporal_state_gpu,
            temporal_state_py,
            msg="GPU temporal state should match Python",
        )
        expected_accepted = torch.tensor(
            input_batch_py.num_accepted_tokens_cpu[:num_reqs],
            dtype=torch.int32,
            device=device,
        )
        torch.testing.assert_close(
            gpu_ctx.num_accepted_tokens_out[:num_reqs],
            expected_accepted,
            msg="GPU num_accepted_tokens should match Python",
        )

    def test_different_block_idx_copies_without_setting_accepted_to_1(
        self, device, test_config
    ):
        """
        Test that neither special-case path triggers when
        src_block_idx != dest_block_idx, and GPU matches Python behavior.

        When copying between different blocks:
        1. src_addr != dst_addr (different blocks = different addresses)
        2. src_block_idx != dest_block_idx

        Therefore:
        - The ``src_addr == dst_addr`` early-return does NOT trigger
        - The ``src_block_idx == dest_block_idx`` post-copy update does NOT trigger
        - Copy happens normally
        - num_accepted_tokens remains UNCHANGED

        Test setup (block_size=16):
        - num_tokens_running_state = 60 + 3 - 0 = 63
        - new_num_computed = 63 + 3 - 1 = 65
        - aligned_new_computed = 64
        - accept_token_bias = 64 - 63 = 1
        - dest_block_idx = 64 // 16 - 1 = 3
        - src_block_idx = 2 (set explicitly, != dest_block_idx)
        """
        cfg = test_config
        torch.manual_seed(1003)

        req_ids = ["req_0"]
        num_computed_tokens = [60]
        num_scheduled_tokens = {"req_0": 3}
        num_draft_tokens: dict[str, int] = {}
        num_accepted_tokens = [3]  # Should remain 3, NOT set to 1
        mamba_state_idx = [2]  # src_block_idx = 2, dest_block_idx will be 3
        block_ids_per_req = [list(range(8))]

        layer_names = ["layer_0"]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        (
            conv_state_py,
            temporal_state_py,
            conv_state_gpu,
            temporal_state_gpu,
            forward_context_py,
            forward_context_gpu,
        ) = _make_dual_layer_state(cfg, device)

        # Clone to verify modification
        conv_state_orig = conv_state_py.clone()

        # --- Run Python path ---
        scheduler_output = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        input_batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)

        postprocess_mamba(
            scheduler_output,
            kv_cache_config,
            input_batch_py,
            requests,
            forward_context_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # --- Run GPU path ---
        gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)

        num_reqs = len(req_ids)
        block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
        block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # --- Verify Python behavior (ground truth) ---
        dest_block_id = block_ids_per_req[0][3]  # dest_block_idx = 3

        # Copy DID happen (dest block should be modified)
        conv_changed = not torch.allclose(
            conv_state_py[dest_block_id], conv_state_orig[dest_block_id]
        )
        assert conv_changed, "Python: Conv state copy should have happened"

        # num_accepted_tokens should NOT be changed to 1
        assert input_batch_py.num_accepted_tokens_cpu[0] == num_accepted_tokens[0], (
            f"Python: num_accepted_tokens should remain {num_accepted_tokens[0]}, "
            f"got {input_batch_py.num_accepted_tokens_cpu[0]}"
        )

        # --- Verify GPU matches Python ---
        torch.testing.assert_close(
            conv_state_gpu,
            conv_state_py,
            msg="GPU conv state should match Python",
        )
        torch.testing.assert_close(
            temporal_state_gpu,
            temporal_state_py,
            msg="GPU temporal state should match Python",
        )
        expected_accepted = torch.tensor(
            input_batch_py.num_accepted_tokens_cpu[:num_reqs],
            dtype=torch.int32,
            device=device,
        )
        torch.testing.assert_close(
            gpu_ctx.num_accepted_tokens_out[:num_reqs],
            expected_accepted,
            msg="GPU num_accepted_tokens should match Python",
        )

    def test_prefix_caching_shared_block_does_not_set_accepted_to_1(
        self, device, test_config
    ):
        """
        Regression test: with prefix caching, different logical block indices
        can map to the same physical block. The kernel must NOT set
        num_accepted_tokens to 1 in that case.

        When src_block_idx != dest_block_idx but block_table maps both to the
        same physical block ID, src_addr == dst_addr. The copy is correctly
        skipped (self-copy is a no-op), but num_accepted_tokens must be
        preserved — only logical-index equality justifies setting it to 1.

        Test setup (block_size=16):
        - num_tokens_running_state = 30 + 2 - 0 = 32
        - new_num_computed = 32 + 3 - 1 = 34
        - aligned_new_computed = 32
        - accept_token_bias = 32 - 32 = 0
        - dest_block_idx = 32 // 16 - 1 = 1
        - src_block_idx = 0 (set explicitly, != dest_block_idx)
        - block_ids = [5, 5, ...] — prefix caching: both logical indices
          map to the same physical block 5
        """
        cfg = test_config
        torch.manual_seed(2001)

        req_ids = ["req_0"]
        num_computed_tokens = [30]
        num_scheduled_tokens = {"req_0": 2}
        num_draft_tokens: dict[str, int] = {}
        num_accepted_tokens = [3]  # Must stay 3, NOT become 1
        mamba_state_idx = [0]  # src_block_idx = 0, dest_block_idx will be 1

        # Prefix caching: logical blocks 0 and 1 share physical block 5
        block_ids_per_req = [[5, 5, 2, 3, 4, 6, 7, 8]]

        layer_names = ["layer_0"]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        (
            conv_state_py,
            temporal_state_py,
            conv_state_gpu,
            temporal_state_gpu,
            forward_context_py,
            forward_context_gpu,
        ) = _make_dual_layer_state(cfg, device)

        conv_state_orig = conv_state_py.clone()
        temporal_state_orig = temporal_state_py.clone()

        # --- Run Python path ---
        scheduler_output = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        input_batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)

        postprocess_mamba(
            scheduler_output,
            kv_cache_config,
            input_batch_py,
            requests,
            forward_context_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # --- Run GPU path ---
        gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
        num_reqs = len(req_ids)
        block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
        block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # --- Verify Python behavior (ground truth) ---
        # Copy is self-to-self (same physical block), state unchanged
        torch.testing.assert_close(
            conv_state_py,
            conv_state_orig,
            msg="Python: Conv state should be unchanged (self-copy)",
        )
        torch.testing.assert_close(
            temporal_state_py,
            temporal_state_orig,
            msg="Python: Temporal state should be unchanged (self-copy)",
        )
        # num_accepted_tokens must NOT be set to 1 (src_block_idx != dest_block_idx)
        assert input_batch_py.num_accepted_tokens_cpu[0] == num_accepted_tokens[0], (
            f"Python: num_accepted_tokens should remain {num_accepted_tokens[0]}, "
            f"got {input_batch_py.num_accepted_tokens_cpu[0]}"
        )

        # --- Verify GPU matches Python ---
        torch.testing.assert_close(
            conv_state_gpu,
            conv_state_py,
            msg="GPU conv state should match Python",
        )
        torch.testing.assert_close(
            temporal_state_gpu,
            temporal_state_py,
            msg="GPU temporal state should match Python",
        )
        expected_accepted = torch.tensor(
            input_batch_py.num_accepted_tokens_cpu[:num_reqs],
            dtype=torch.int32,
            device=device,
        )
        torch.testing.assert_close(
            gpu_ctx.num_accepted_tokens_out[:num_reqs],
            expected_accepted,
            msg="GPU num_accepted_tokens should match Python (must NOT be 1)",
        )

    def test_prefix_caching_nonsequential_block_ids_boundary(self, device, test_config):
        """
        Regression test: non-sequential physical block IDs under prefix caching
        with the needs_copy boundary at exact equality.

        Under PC, the block allocator assigns physical block IDs in arbitrary
        order (e.g., [17, 3, 42, 9] instead of [0, 1, 2, 3]). The needs_copy
        condition is purely token-count based and must evaluate identically
        regardless of the physical block IDs assigned. This test verifies that
        the kernel's address arithmetic (block_table lookup, stride computation)
        produces correct copies when physical IDs are non-sequential.

        Two requests exercise different boundary behaviors:
        - req_0: aligned_new_computed == num_tokens_running_state (exact boundary)
          This is the tightest edge: one fewer accepted token and no copy needed.
        - req_1: aligned_new_computed == num_tokens_running_state (exact boundary)
          Different block layout, src!=dest, real copy happens.

        Both use non-sequential block IDs typical of PC reuse patterns.

        Test setup (block_size=16):
        req_0:
        - num_tokens_running_state = 48 + 0 - 0 = 48
        - new_num_computed = 48 + 1 - 1 = 48
        - aligned_new_computed = 48
        - needs_copy = (48 >= 48) = True (exact boundary!)
        - accept_token_bias = 48 - 48 = 0
        - dest_block_idx = 48 // 16 - 1 = 2
        - src_block_idx = 2 (same as dest -> num_accepted = 1)

        req_1:
        - num_tokens_running_state = 31 + 1 - 0 = 32
        - new_num_computed = 32 + 3 - 1 = 34
        - aligned_new_computed = 32
        - needs_copy = (32 >= 32) = True (exact boundary!)
        - accept_token_bias = 32 - 32 = 0
        - dest_block_idx = 32 // 16 - 1 = 1
        - src_block_idx = 0 (diff from dest -> num_accepted unchanged)
        """
        cfg = test_config
        torch.manual_seed(4001)

        req_ids = ["req_0", "req_1"]
        num_computed_tokens = [48, 31]
        num_scheduled_tokens = {"req_0": 0, "req_1": 1}
        num_draft_tokens: dict[str, int] = {}
        num_accepted_tokens = [1, 3]
        mamba_state_idx = [2, 0]

        # Non-sequential block IDs typical of prefix caching allocation
        block_ids_per_req = [
            [17, 3, 42, 9, 25, 11, 30, 2],  # req_0: scattered physical blocks
            [41, 7, 22, 15, 38, 19, 4, 28],  # req_1: different scattered blocks
        ]

        layer_names = [f"layer_{i}" for i in range(cfg.num_layers)]
        # Need enough physical blocks for the scattered IDs
        num_blocks = 50
        local_cfg = _TestConfig(num_blocks=num_blocks, max_num_reqs=cfg.max_num_reqs)
        kv_cache_config = _make_kv_cache_config(local_cfg, layer_names)

        (
            conv_states_py,
            temporal_states_py,
            conv_states_gpu,
            temporal_states_gpu,
            forward_context_py,
            forward_context_gpu,
        ) = _make_dual_states(local_cfg, layer_names, device)

        # --- Run Python path ---
        scheduler_output = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        input_batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(local_cfg, kv_cache_config, device)

        postprocess_mamba(
            scheduler_output,
            kv_cache_config,
            input_batch_py,
            requests,
            forward_context_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # --- Run GPU path ---
        gpu_ctx = _make_gpu_ctx(local_cfg, kv_cache_config, device)
        num_reqs = len(req_ids)
        max_blocks = max(len(b) for b in block_ids_per_req)
        block_table_gpu = torch.zeros(
            num_reqs, max_blocks, dtype=torch.int32, device=device
        )
        for i, block_ids in enumerate(block_ids_per_req):
            block_table_gpu[i, : len(block_ids)] = torch.tensor(
                block_ids, dtype=torch.int32
            )

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # --- Compare results ---
        for i in range(cfg.num_layers):
            torch.testing.assert_close(
                conv_states_gpu[i],
                conv_states_py[i],
                msg=f"Conv state mismatch at layer {i} with non-sequential block IDs",
            )
            torch.testing.assert_close(
                temporal_states_gpu[i],
                temporal_states_py[i],
                msg=(
                    f"Temporal state mismatch at layer {i} "
                    f"with non-sequential block IDs"
                ),
            )

        expected_accepted = torch.tensor(
            input_batch_py.num_accepted_tokens_cpu[:num_reqs],
            dtype=torch.int32,
            device=device,
        )
        torch.testing.assert_close(
            gpu_ctx.num_accepted_tokens_out[:num_reqs],
            expected_accepted,
            msg="num_accepted_tokens mismatch with non-sequential block IDs",
        )

        # Verify req_0 had num_accepted set to 1 (src==dest) and req_1 unchanged
        assert input_batch_py.num_accepted_tokens_cpu[0] == 1
        assert input_batch_py.num_accepted_tokens_cpu[1] == num_accepted_tokens[1]

    def test_prefix_caching_mixed_shared_and_distinct_blocks(self, device, test_config):
        """
        Regression test: mixed batch under prefix caching where some requests
        have shared physical blocks (aliased) and others have distinct blocks,
        with the needs_copy boundary at various positions.

        This tests the interaction between:
        1. PC block aliasing (src and dest map to same physical block)
        2. The needs_copy boundary (exact equality vs well-past vs no-copy)
        3. Non-sequential physical block IDs

        Batch of 4 requests:
        - req_0: needs_copy=True, src!=dest, shared physical block (PC aliased)
                 -> copy skipped (src_addr==dst_addr), num_accepted PRESERVED
        - req_1: needs_copy=True, src!=dest, distinct blocks, non-sequential IDs
                 -> real copy happens, num_accepted PRESERVED
        - req_2: needs_copy=False (below boundary)
                 -> no action at all
        - req_3: needs_copy=True, src==dest (exact boundary, zero bias)
                 -> copy skipped (self-copy), num_accepted SET TO 1

        Test setup (block_size=16):
        req_0: running=30+2-0=32, new=32+3-1=34, aligned=32, 32>=32 -> COPY
               bias=0, dest=32//16-1=1, src=0 (!=dest)
               block_ids=[5,5,...] -> same physical -> skip, keep accepted=3

        req_1: running=60+5-2=63, new=63+3-1=65, aligned=64, 64>=63 -> COPY
               bias=1, dest=64//16-1=3, src=2 (!=dest)
               block_ids=[41,7,22,15,...] -> distinct -> real copy, keep accepted=3

        req_2: running=30+3-0=33, new=33+1-1=33, aligned=32, 32<33 -> NO COPY

        req_3: running=48+0-0=48, new=48+1-1=48, aligned=48, 48>=48 -> COPY
               bias=0, dest=48//16-1=2, src=2 (==dest)
               block_ids=[10,20,30,...] -> distinct IDs, same logical idx
               -> self-copy (src_addr==dst_addr), set accepted=1
        """
        cfg = test_config
        torch.manual_seed(5001)

        req_ids = ["req_0", "req_1", "req_2", "req_3"]
        num_computed_tokens = [30, 60, 30, 48]
        num_scheduled_tokens = {"req_0": 2, "req_1": 5, "req_2": 3, "req_3": 0}
        num_draft_tokens = {"req_1": 2}
        num_accepted_tokens = [3, 3, 1, 1]
        mamba_state_idx = [0, 2, 1, 2]

        # Block IDs with various PC patterns:
        # req_0: shared blocks (PC alias: logical 0 and 1 -> physical 5)
        # req_1: distinct non-sequential blocks
        # req_2: doesn't matter (no copy)
        # req_3: distinct sequential blocks (no aliasing)
        block_ids_per_req = [
            [5, 5, 12, 18, 23, 31, 44, 2],  # req_0: blocks 0,1 share phys 5
            [41, 7, 22, 15, 38, 19, 4, 28],  # req_1: all distinct
            [10, 20, 30, 40, 1, 6, 8, 14],  # req_2: irrelevant
            [10, 20, 30, 40, 1, 6, 8, 14],  # req_3: distinct, dest=src=idx 2
        ]

        layer_names = [f"layer_{i}" for i in range(cfg.num_layers)]
        num_blocks = 50
        local_cfg = _TestConfig(num_blocks=num_blocks, max_num_reqs=cfg.max_num_reqs)
        kv_cache_config = _make_kv_cache_config(local_cfg, layer_names)

        (
            conv_states_py,
            temporal_states_py,
            conv_states_gpu,
            temporal_states_gpu,
            forward_context_py,
            forward_context_gpu,
        ) = _make_dual_states(local_cfg, layer_names, device)

        # --- Run Python path ---
        scheduler_output = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        input_batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(local_cfg, kv_cache_config, device)

        postprocess_mamba(
            scheduler_output,
            kv_cache_config,
            input_batch_py,
            requests,
            forward_context_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # --- Run GPU path ---
        gpu_ctx = _make_gpu_ctx(local_cfg, kv_cache_config, device)
        num_reqs = len(req_ids)
        max_blocks = max(len(b) for b in block_ids_per_req)
        block_table_gpu = torch.zeros(
            num_reqs, max_blocks, dtype=torch.int32, device=device
        )
        for i, block_ids in enumerate(block_ids_per_req):
            block_table_gpu[i, : len(block_ids)] = torch.tensor(
                block_ids, dtype=torch.int32
            )

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # --- Compare all state tensors ---
        for i in range(cfg.num_layers):
            torch.testing.assert_close(
                conv_states_gpu[i],
                conv_states_py[i],
                msg=(
                    f"Conv state mismatch at layer {i} — "
                    f"mixed PC batch with shared/distinct blocks"
                ),
            )
            torch.testing.assert_close(
                temporal_states_gpu[i],
                temporal_states_py[i],
                msg=(
                    f"Temporal state mismatch at layer {i} — "
                    f"mixed PC batch with shared/distinct blocks"
                ),
            )

        # --- Compare num_accepted_tokens ---
        expected_accepted = torch.tensor(
            input_batch_py.num_accepted_tokens_cpu[:num_reqs],
            dtype=torch.int32,
            device=device,
        )
        torch.testing.assert_close(
            gpu_ctx.num_accepted_tokens_out[:num_reqs],
            expected_accepted,
            msg="num_accepted_tokens mismatch in mixed PC batch",
        )

        # Verify per-request expectations:
        # req_0: src!=dest, shared block -> preserved (3)
        assert input_batch_py.num_accepted_tokens_cpu[0] == 3
        # req_1: src!=dest, distinct blocks -> preserved (3)
        assert input_batch_py.num_accepted_tokens_cpu[1] == 3
        # req_2: no copy -> preserved (1)
        assert input_batch_py.num_accepted_tokens_cpu[2] == 1
        # req_3: src==dest -> set to 1
        assert input_batch_py.num_accepted_tokens_cpu[3] == 1

    def test_pc_aliased_blocks_skip_must_use_logical_idx_not_addr(
        self, device, test_config
    ):
        """
        Regression test for 6466ce0d vs 959ca0fd: the kernel's early-return
        guard must compare logical block indices, not physical addresses.

        Under prefix caching, different logical blocks (src_block_idx=0,
        dest_block_idx=1) can map to the same physical block. When
        accept_token_bias=0, this makes src_addr == dst_addr for BOTH conv
        and temporal states. A buggy guard `if src_addr == dst_addr` would
        incorrectly set num_accepted_tokens=1; the correct guard is
        `if src_block_idx == dest_block_idx and accept_token_bias == 0`.

        The Python reference only sets num_accepted_tokens=1 when
        src_block_idx == dest_block_idx (line 79 of postprocess_mamba).
        With src_block_idx=0, dest_block_idx=1, num_accepted_tokens must
        be preserved even though the physical addresses match.

        Test setup (block_size=16):
        - num_tokens_running_state = 30 + 2 - 0 = 32
        - new_num_computed = 32 + 3 - 1 = 34
        - aligned_new_computed = 32
        - needs_copy = (32 >= 32) = True
        - accept_token_bias = 32 - 32 = 0
        - dest_block_idx = 32 // 16 - 1 = 1
        - src_block_idx = 0 (explicitly, != dest_block_idx)
        - block_ids = [7, 7, ...] -> physical aliasing via prefix caching

        Expected: num_accepted_tokens stays 3 (not set to 1).
        Bug (959ca0fd): kernel saw src_addr == dst_addr, set it to 1.
        """
        cfg = test_config
        torch.manual_seed(6001)

        req_ids = ["req_0"]
        num_computed_tokens = [30]
        num_scheduled_tokens = {"req_0": 2}
        num_draft_tokens: dict[str, int] = {}
        num_accepted_tokens = [3]
        mamba_state_idx = [0]  # src_block_idx = 0

        # Prefix caching: logical blocks 0 and 1 both map to physical block 7.
        block_ids_per_req = [[7, 7, 10, 11, 12, 13, 14, 15]]

        layer_names = ["layer_0"]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        (
            conv_state_py,
            temporal_state_py,
            conv_state_gpu,
            temporal_state_gpu,
            forward_context_py,
            forward_context_gpu,
        ) = _make_dual_layer_state(cfg, device)

        # --- Run Python path ---
        scheduler_output = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        input_batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)

        postprocess_mamba(
            scheduler_output,
            kv_cache_config,
            input_batch_py,
            requests,
            forward_context_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # Python reference: src_block_idx(0) != dest_block_idx(1) -> no change
        assert input_batch_py.num_accepted_tokens_cpu[0] == 3, (
            f"Python: num_accepted_tokens should remain 3, "
            f"got {input_batch_py.num_accepted_tokens_cpu[0]}"
        )

        # --- Run GPU path ---
        gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
        num_reqs = len(req_ids)
        block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
        block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # The critical assertion: kernel must NOT set num_accepted_tokens to 1
        # when src_block_idx != dest_block_idx, even though src_addr == dst_addr
        # due to prefix caching aliasing.
        #
        # Old kernel (959ca0fd): `if src_addr == dst_addr` -> FAILS here (sets 1)
        # Fixed kernel (6466ce0d): `if src_block_idx == dest_block_idx and
        #   accept_token_bias == 0` -> PASSES (preserves 3)
        kernel_accepted = gpu_ctx.num_accepted_tokens_out[0].item()
        assert kernel_accepted == 3, (
            f"Kernel set num_accepted_tokens to {kernel_accepted} but expected 3. "
            f"The early-return guard likely compared physical addresses "
            f"(src_addr == dst_addr) instead of logical block indices "
            f"(src_block_idx == dest_block_idx). Under prefix caching, "
            f"different logical blocks can share the same physical block."
        )

        # Also verify state tensors match Python
        torch.testing.assert_close(
            conv_state_gpu,
            conv_state_py,
            msg="GPU conv state should match Python",
        )
        torch.testing.assert_close(
            temporal_state_gpu,
            temporal_state_py,
            msg="GPU temporal state should match Python",
        )

    def test_as_strided_temporal_copy_size(self, device, test_config):
        """
        Regression test for 240723d46: temporal copy_size must be
        inner_size * elem_size, not state_block_stride.

        In production (gpu_model_runner.py), conv and temporal states share
        a raw buffer via torch.as_strided where stride(0) equals
        page_size_bytes / elem_size — larger than either state's natural
        element count.  Using stride(0) as copy_size for temporal states
        overwrites into the next block's conv region.

        Layout per page (384 float16 elements = 768 bytes):
            [conv: 256 elems | temporal: 128 elems]

        The test triggers a temporal copy from block 4 to block 3.  With the
        bug the kernel copies 768 bytes (page stride) instead of 256 bytes
        (128 * 2), overwriting conv_state[4] with conv_state[5]'s data.

        Test setup (block_size=16):
        - running = 60 + 5 - 2 = 63
        - new = 63 + 3 - 1 = 65
        - aligned = 64 >= 63 -> COPY needed
        - accept_token_bias = 64 - 63 = 1
        - dest_block_idx = 64 // 16 - 1 = 3
        - temporal: actual_src_block_idx = 3 + 1 = 4  (block_ids[4] = 4)
        """
        cfg = test_config
        torch.manual_seed(7001)

        req_ids = ["req_0"]
        num_computed_tokens = [60]
        num_scheduled_tokens = {"req_0": 5}
        num_draft_tokens = {"req_0": 2}
        num_accepted_tokens = [3]
        mamba_state_idx = [3]
        block_ids_per_req = [list(range(8))]

        layer_names = ["layer_0"]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        # --- Production-like packed layout (mirrors gpu_model_runner.py) ---
        conv_shape = (cfg.conv_width, cfg.conv_inner_dim)
        temporal_shape = (cfg.temporal_state_dim,)
        dtype = cfg.dtype
        elem_size = torch.tensor([], dtype=dtype).element_size()

        conv_natural_elems = cfg.conv_width * cfg.conv_inner_dim
        temporal_natural_elems = cfg.temporal_state_dim
        page_size_bytes = (conv_natural_elems + temporal_natural_elems) * elem_size
        num_element_per_page = page_size_bytes // elem_size

        assert num_element_per_page > temporal_natural_elems, (
            "Test requires padded stride; page must be larger than one state"
        )

        raw_py = torch.randn(
            cfg.num_blocks * num_element_per_page, dtype=dtype, device=device
        )
        raw_gpu = raw_py.clone()

        def make_views(raw):
            conv_tgt = (cfg.num_blocks, *conv_shape)
            conv_nat_stride = torch.empty(conv_tgt).stride()
            conv = torch.as_strided(
                raw,
                size=conv_tgt,
                stride=(num_element_per_page, *conv_nat_stride[1:]),
                storage_offset=0,
            )

            temp_tgt = (cfg.num_blocks, *temporal_shape)
            temp_nat_stride = torch.empty(temp_tgt).stride()
            temp = torch.as_strided(
                raw,
                size=temp_tgt,
                stride=(num_element_per_page, *temp_nat_stride[1:]),
                storage_offset=conv_natural_elems,
            )
            return conv, temp

        conv_py, temp_py = make_views(raw_py)
        conv_gpu, temp_gpu = make_views(raw_gpu)

        fwd_py = {"layer_0": _make_mock_attention(conv_py, temp_py)}
        fwd_gpu = {"layer_0": _make_mock_attention(conv_gpu, temp_gpu)}

        # --- Python reference ---
        sched = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)

        postprocess_mamba(
            sched,
            kv_cache_config,
            batch_py,
            requests,
            fwd_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # --- GPU fused kernel ---
        gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
        num_reqs = 1
        block_table = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
        block_table[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=fwd_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # --- Assertions ---
        # With the bug (pre-240723d46), the kernel copies page_size_bytes
        # (768) for temporal state instead of 256 bytes, overwriting
        # conv_state[4] with conv_state[5]'s data.
        torch.testing.assert_close(
            conv_gpu,
            conv_py,
            msg=(
                "Conv state corrupted: temporal copy_size was likely "
                "state_block_stride instead of inner_size * elem_size"
            ),
        )
        torch.testing.assert_close(
            temp_gpu,
            temp_py,
            msg="Temporal state mismatch",
        )

        expected_accepted = torch.tensor(
            batch_py.num_accepted_tokens_cpu[:num_reqs],
            dtype=torch.int32,
            device=device,
        )
        torch.testing.assert_close(
            gpu_ctx.num_accepted_tokens_out[:num_reqs],
            expected_accepted,
            msg="num_accepted_tokens mismatch",
        )

    def test_temporal_copy_with_bias_ge_2(self, device, test_config):
        """
        Coverage test for the temporal-state block-table stride arithmetic
        when ``accept_token_bias >= 2``.

        The kernel computes, for temporal (non-conv) states::

            actual_src_block_idx = src_block_idx + accept_token_bias
            actual_src_block_id = block_table[req, actual_src_block_idx]

        All prior regression tests exercise only ``bias == 1``, i.e. they
        only ever read one slot ahead of ``src_block_idx`` in the block
        table. An off-by-one (or missing scale) in the address computation
        on line 143 of ``mamba_utils.py`` would be invisible to every
        existing test but would silently read the wrong physical block on
        any speculative-decode cycle that accepts multiple tokens across a
        block boundary, feeding a stale hidden state forward one step.

        Setup (block_size=16):
        - running   = 28 + 2 - 0 = 30
        - new       = 30 + 3 - 1 = 32
        - aligned   = 32 >= 30 -> COPY needed
        - bias      = 32 - 30 = 2             (key: >= 2)
        - dest_idx  = 32 // 16 - 1 = 1
        - src_idx   = 1 (same as dest -> exercises post-copy accepted=1 write)
        - temporal actual_src_block_idx = 1 + 2 = 3 (reads block_table[0, 3])

        With identity block_ids = [0,1,2,3,...], an off-by-one that used
        bias=1 would copy from block_ids[2]=2 instead of block_ids[3]=3,
        producing a clear state-value mismatch against the Python
        reference.
        """
        cfg = test_config
        torch.manual_seed(7002)

        req_ids = ["req_0"]
        num_computed_tokens = [28]
        num_scheduled_tokens = {"req_0": 2}
        num_draft_tokens: dict[str, int] = {}
        num_accepted_tokens = [3]  # -> accept_token_bias = 2
        mamba_state_idx = [1]  # src_block_idx = 1 = dest_block_idx
        block_ids_per_req = [list(range(8))]

        layer_names = ["layer_0"]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        (
            conv_state_py,
            temporal_state_py,
            conv_state_gpu,
            temporal_state_gpu,
            fwd_py,
            fwd_gpu,
        ) = _make_dual_layer_state(cfg, device)
        temporal_state_orig = temporal_state_py.clone()

        # --- Python reference ---
        sched = _make_postprocess_scheduler_output(
            req_ids,
            num_scheduled_tokens,
            {k: [None] * v for k, v in num_draft_tokens.items() if v > 0},
        )
        batch_py = _make_input_batch(
            req_ids, num_accepted_tokens.copy(), mamba_state_idx.copy()
        )
        requests = _make_requests(req_ids, num_computed_tokens, block_ids_per_req)
        copy_bufs = _make_copy_bufs(cfg, kv_cache_config, device)

        postprocess_mamba(
            sched,
            kv_cache_config,
            batch_py,
            requests,
            fwd_py,
            _COPY_FUNCS,
            copy_bufs,
        )
        torch.accelerator.synchronize()

        # --- GPU fused kernel ---
        gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device)
        num_reqs = 1
        block_table = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
        block_table[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)

        _run_gpu_postprocess(
            gpu_ctx,
            kv_cache_config=kv_cache_config,
            forward_context=fwd_gpu,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )

        # --- Ground truth: Python must have sourced temporal from block 3 ---
        actual_src_block_id = block_ids_per_req[0][3]  # == 3
        dest_block_id = block_ids_per_req[0][1]  # == 1
        torch.testing.assert_close(
            temporal_state_py[dest_block_id],
            temporal_state_orig[actual_src_block_id],
            msg=(
                "Python reference did not copy from block_ids[src+bias]=3; "
                "test preconditions are wrong"
            ),
        )

        # --- GPU kernel must match Python byte-for-byte ---
        torch.testing.assert_close(
            conv_state_gpu,
            conv_state_py,
            msg="Conv state mismatch at accept_token_bias=2",
        )
        torch.testing.assert_close(
            temporal_state_gpu,
            temporal_state_py,
            msg=(
                "Temporal state mismatch at accept_token_bias=2: the kernel "
                "likely read the wrong slot of the block table "
                "(actual_src_block_idx stride arithmetic)"
            ),
        )

        expected_accepted = torch.tensor(
            batch_py.num_accepted_tokens_cpu[:num_reqs],
            dtype=torch.int32,
            device=device,
        )
        torch.testing.assert_close(
            gpu_ctx.num_accepted_tokens_out[:num_reqs],
            expected_accepted,
            msg="num_accepted_tokens mismatch at accept_token_bias=2",
        )

    def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd(
        self, device, test_config, monkeypatch
    ):
        """DS conv postprocess should match SD when accept_token_bias > 0."""
        from vllm.model_executor.layers.mamba import mamba_utils as model_mamba_utils

        cfg = test_config
        torch.manual_seed(38898)

        req_ids = ["req_0"]
        num_computed_tokens = [30]
        num_scheduled_tokens = {"req_0": 1}
        num_draft_tokens: dict[str, int] = {}
        num_accepted_tokens = [2]  # Results in accept_token_bias = 1
        mamba_state_idx = [1]  # src_block_idx = 1 = dest_block_idx
        block_ids_per_req = [list(range(8))]

        layer_names = ["layer_0"]
        kv_cache_config = _make_kv_cache_config(cfg, layer_names)

        num_reqs = len(req_ids)
        block_table_gpu = torch.zeros(num_reqs, 8, dtype=torch.int32, device=device)
        block_table_gpu[0, :8] = torch.tensor(block_ids_per_req[0], dtype=torch.int32)

        # Same logical conv state in SD and DS layouts.
        sd_source_conv = torch.randn(
            cfg.num_blocks,
            cfg.conv_width,
            cfg.conv_inner_dim,
            dtype=cfg.dtype,
            device=device,
        )
        ds_source_conv = sd_source_conv.permute(0, 2, 1).contiguous()
        sd_source_temporal = torch.randn(
            cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device
        )

        # SD GPU path. Default layout is SD.
        model_mamba_utils.get_conv_state_layout.cache_clear()
        sd_conv = sd_source_conv.clone()
        sd_temporal = sd_source_temporal.clone()
        forward_context_sd = {
            "layer_0": _make_mock_attention(sd_conv, sd_temporal),
        }
        gpu_ctx_sd = _make_gpu_ctx(cfg, kv_cache_config, device)
        _run_gpu_postprocess(
            gpu_ctx_sd,
            kv_cache_config=kv_cache_config,
            forward_context=forward_context_sd,
            copy_funcs=_COPY_FUNCS,
            block_table=block_table_gpu,
            req_ids=req_ids,
            num_accepted_tokens=num_accepted_tokens,
            mamba_state_idx=mamba_state_idx,
            num_scheduled_tokens=num_scheduled_tokens,
            num_computed_tokens=num_computed_tokens,
            num_draft_tokens=num_draft_tokens,
            device=device,
        )
        torch.accelerator.synchronize()

        # Sanity: SD path actually modified the state (copy was performed).
        assert not torch.equal(sd_conv, sd_source_conv), (
            "SD baseline did not modify conv state; test setup is wrong"
        )

        # DS GPU path on the DS twin.
        monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", "DS")
        model_mamba_utils.get_conv_state_layout.cache_clear()
        try:
            ds_conv = ds_source_conv.clone()
            ds_temporal = sd_source_temporal.clone()
            forward_context_ds = {
                "layer_0": _make_mock_attention(ds_conv, ds_temporal),
            }
            gpu_ctx_ds = _make_gpu_ctx(cfg, kv_cache_config, device)
            _run_gpu_postprocess(
                gpu_ctx_ds,
                kv_cache_config=kv_cache_config,
                forward_context=forward_context_ds,
                copy_funcs=_COPY_FUNCS,
                block_table=block_table_gpu,
                req_ids=req_ids,
                num_accepted_tokens=num_accepted_tokens,
                mamba_state_idx=mamba_state_idx,
                num_scheduled_tokens=num_scheduled_tokens,
                num_computed_tokens=num_computed_tokens,
                num_draft_tokens=num_draft_tokens,
                device=device,
            )
            torch.accelerator.synchronize()
        finally:
            # Reset the lru cache so other tests see the default layout again.
            model_mamba_utils.get_conv_state_layout.cache_clear()

        # DS bytes, un-permuted, should match the SD result.
        torch.testing.assert_close(
            ds_conv.permute(0, 2, 1).contiguous(),
            sd_conv,
            msg=(
                "DS conv post-kernel does not match SD baseline; the DS "
                "row-loop in postprocess_mamba_fused_kernel is wrong."
            ),
        )
        torch.testing.assert_close(
            ds_temporal,
            sd_temporal,
            msg="DS temporal state diverged from SD",
        )
        torch.testing.assert_close(
            gpu_ctx_ds.num_accepted_tokens_out[:num_reqs],
            gpu_ctx_sd.num_accepted_tokens_out[:num_reqs],
            msg="DS num_accepted_tokens diverged from SD",
        )
