wehub-resource-sync
2250 lines
82 KiB
Python
2250 lines
82 KiB
Python
# 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",
|
|
)
|