minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[KDA] Add target_verify support for speculative decoding (#26888)

Browse Source 
Co-authored-by: yuyanqi <yuyanqi@meituan.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
This commit is contained in:
yyqjwyy
2026-07-25 19:52:44 +08:00
committed by GitHub
co-authored by yuyanqi Claude Opus 4.6 Xinyuan Tong
parent 7c4b22fae5
commit a678a42033
6 changed files with 701 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/linear/kda_backend.py
+5
View File
@@ -403,6 +403,11 @@ class KDAAttnBackend(MambaAttnBackendBase):
ssm_states = mamba_cache_params.temporal
# Normal extend path
if forward_batch.extend_prefix_lens is None:
raise RuntimeError(
"extend_prefix_lens cannot be None in non-TARGET_VERIFY mode."
)
has_initial_state = forward_batch.extend_prefix_lens > 0
if self.forward_metadata.has_mamba_track_mask:
python/sglang/srt/mem_cache/memory_pool.py
+32
View File
@@ -381,6 +381,34 @@ class MambaPool:
intermediate_ssm: Optional[torch.Tensor]
intermediate_conv_window: List[torch.Tensor]
def _detect_conv_window_axis(
self, conv_state_shape: List[Tuple[int, int]], win_len: int
) -> int:
"""Prefer GDN's trailing axis when both match; mixed layer layouts cannot
share one overlapping conv-window buffer.
"""
axis = None
for conv_shape in conv_state_shape:
if conv_shape[-1] == win_len:
shape_axis = len(conv_shape) - 1
elif conv_shape[0] == win_len:
shape_axis = 0
else:
raise ValueError(
f"conv_state shape {conv_shape} has no axis of length "
f"conv_kernel-1={win_len}; cannot build the deduplicated "
"sliding-window conv-intermediate view."
)
if axis is None:
axis = shape_axis
elif axis != shape_axis:
raise ValueError(
"inconsistent conv-window axis across conv shapes "
f"{conv_state_shape}; a single conv_window_axis cannot serve "
"mixed layouts."
)
return axis
def _allocate_deduplicated_conv_window(
self,
*,
@@ -676,6 +704,10 @@ class MambaPool:
)
self._intermediate_conv_window_phys = []
if dedup_conv_window:
win_len = cache_params.shape.conv_kernel - 1
self.conv_window_axis = self._detect_conv_window_axis(
conv_state_shape, win_len
)
intermediate_conv_window_cache = []
for conv_shape in conv_state_shape:
phys, view = self._allocate_deduplicated_conv_window(
python/sglang/srt/models/kimi_linear.py
+6 -4
View File
@@ -382,10 +382,12 @@ class KimiDeltaAttention(nn.Module):
hidden_states
)
# For prefill: raw gate is passed to chunk_kda_fwd, which fuses gate
# activation with chunk_local_cumsum (kda_gate_chunk_cumsum kernel).
# For decode: gate activation is handled inside fused_recurrent kernel.
if not forward_batch.forward_mode.is_decode():
# Prefill passes raw gates to chunk KDA; decode and target-verify kernels
# apply the activation internally.
if (
not forward_batch.forward_mode.is_decode()
and not forward_batch.forward_mode.is_target_verify()
):
forget_gate = forget_gate.unflatten(
-1, (-1, self.head_dim)
) # [T, H*K] -> [T, H, K]
test/manual/test_kda_spec_integration.py
+71
View File
@@ -0,0 +1,71 @@
import concurrent.futures
import time
import requests
BASE_URL = "http://localhost:30000"
SHARED_PREFIX = "You are a helpful assistant. " * 20
def test_normal_inference_no_regression():
resp = requests.post(
f"{BASE_URL}/generate",
json={
"text": "What is 2+2?",
"sampling_params": {"max_new_tokens": 32, "temperature": 0.0},
},
)
assert resp.status_code == 200, f"Status {resp.status_code}: {resp.text}"
data = resp.json()
print(f"Normal inference: {data['text'][:80]}")
assert len(data["text"]) > 0
def test_prefix_caching_still_works():
resp1 = requests.post(
f"{BASE_URL}/generate",
json={
"text": SHARED_PREFIX + "What is 1+1?",
"sampling_params": {"max_new_tokens": 32, "temperature": 0.0},
},
)
time.sleep(0.5)
resp2 = requests.post(
f"{BASE_URL}/generate",
json={
"text": SHARED_PREFIX + "What is 3+3?",
"sampling_params": {"max_new_tokens": 32, "temperature": 0.0},
},
)
assert resp1.status_code == 200
assert resp2.status_code == 200
cached = resp2.json().get("meta_info", {}).get("cached_tokens", 0)
print(f"Cached tokens: {cached}")
assert cached > 0, "Prefix caching should work"
def test_batch_inference():
prompts = [f"Count from 1 to {i + 3}" for i in range(8)]
def send(p):
return requests.post(
f"{BASE_URL}/generate",
json={
"text": p,
"sampling_params": {"max_new_tokens": 64, "temperature": 0.0},
},
)
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
futures = [executor.submit(send, p) for p in prompts]
results = [f.result() for f in futures]
for r in results:
assert r.status_code == 200
print(f"Batch test passed: {len(results)} requests OK")
if __name__ == "__main__":
test_normal_inference_no_regression()
test_prefix_caching_still_works()
test_batch_inference()
print("\nAll tests PASSED!")
test/manual/test_kda_target_verify.py
+216
View File
@@ -0,0 +1,216 @@
import torch
def test_kda_target_verify_equivalence():
from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
fused_sigmoid_gating_delta_rule_update,
)
B, HV, K, V = 2, 4, 64, 64
N = 4
device = "cuda"
dtype = torch.float32
torch.manual_seed(42)
q = torch.randn(1, B * N, HV, K, dtype=dtype, device=device)
k = torch.randn(1, B * N, HV, K, dtype=dtype, device=device)
v = torch.randn(1, B * N, HV, V, dtype=dtype, device=device)
a = torch.randn(B * N, HV * K, dtype=dtype, device=device)
b = torch.randn(1, B * N, HV, dtype=dtype, device=device)
A_log = torch.randn(HV, dtype=torch.float32, device=device)
dt_bias = torch.randn(HV, dtype=torch.float32, device=device)
num_slots = B + 2
ssm_states_base = torch.randn(num_slots, HV, K, V, dtype=dtype, device=device)
cache_indices = torch.arange(B, dtype=torch.int32, device=device)
query_start_loc = torch.arange(0, B * N + 1, N, dtype=torch.int32, device=device)
ssm_states_decode = ssm_states_base.clone()
outputs_decode = []
states_after_step = []
for step in range(N):
step_indices = [i * N + step for i in range(B)]
step_q = q[:, step_indices].contiguous()
step_k = k[:, step_indices].contiguous()
step_v = v[:, step_indices].contiguous()
step_a = a[step_indices].contiguous()
step_b = b[:, step_indices].contiguous()
decode_qsl = torch.arange(0, B + 1, dtype=torch.int32, device=device)
out = fused_sigmoid_gating_delta_rule_update(
A_log=A_log,
dt_bias=dt_bias,
q=step_q,
k=step_k,
v=step_v,
a=step_a,
b=step_b,
initial_state_source=ssm_states_decode,
initial_state_indices=cache_indices,
cu_seqlens=decode_qsl,
use_qk_l2norm_in_kernel=True,
softplus_beta=1.0,
softplus_threshold=20.0,
is_kda=True,
)
outputs_decode.append(out)
states_after_step.append(ssm_states_decode[cache_indices].clone())
ssm_states_verify = ssm_states_base.clone()
intermediate_buffer = torch.zeros(
num_slots, N, HV, K, V, dtype=dtype, device=device
)
intermediate_indices = torch.arange(B, dtype=torch.int32, device=device)
out_verify = fused_sigmoid_gating_delta_rule_update(
A_log=A_log,
dt_bias=dt_bias,
q=q,
k=k,
v=v,
a=a,
b=b,
initial_state_source=ssm_states_verify,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=True,
softplus_beta=1.0,
softplus_threshold=20.0,
is_kda=True,
disable_state_update=True,
intermediate_states_buffer=intermediate_buffer,
intermediate_state_indices=intermediate_indices,
cache_steps=N,
retrieve_parent_token=None,
)
out_decode_list = []
for req_idx in range(B):
for step in range(N):
out_decode_list.append(outputs_decode[step][:, req_idx : req_idx + 1])
out_decode_cat = torch.cat(out_decode_list, dim=1)
max_diff = (out_verify - out_decode_cat).abs().max().item()
mean_diff = (out_verify - out_decode_cat).abs().mean().item()
print(f"Output max diff: {max_diff:.6e}, mean diff: {mean_diff:.6e}")
assert max_diff < 1e-5, f"Output mismatch! max diff: {max_diff}"
print("Intermediate state comparison:")
for step in range(N):
for req_idx in range(B):
cached_state = intermediate_buffer[req_idx, step]
decode_state = states_after_step[step][req_idx]
state_diff = (cached_state - decode_state).abs().max().item()
status = "OK" if state_diff < 1e-5 else "FAIL"
print(f" step={step} req={req_idx}: diff={state_diff:.6e} [{status}]")
assert (
state_diff < 1e-5
), f"Intermediate state mismatch at step={step}, req={req_idx}: {state_diff}"
ssm_unchanged_diff = (ssm_states_verify - ssm_states_base).abs().max().item()
print(f"SSM state in-place change (should be 0): {ssm_unchanged_diff:.6e}")
assert (
ssm_unchanged_diff == 0.0
), f"target_verify modified ssm_states in-place! diff: {ssm_unchanged_diff}"
print("\nPASSED: KDA target_verify matches sequential decode!")
def test_kda_target_verify_bf16():
from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
fused_sigmoid_gating_delta_rule_update,
)
B, HV, K, V = 2, 4, 64, 64
N = 4
device = "cuda"
dtype = torch.bfloat16
torch.manual_seed(42)
q = torch.randn(1, B * N, HV, K, dtype=dtype, device=device)
k = torch.randn(1, B * N, HV, K, dtype=dtype, device=device)
v = torch.randn(1, B * N, HV, V, dtype=dtype, device=device)
a = torch.randn(B * N, HV * K, dtype=dtype, device=device)
b = torch.randn(1, B * N, HV, dtype=dtype, device=device)
A_log = torch.randn(HV, dtype=torch.float32, device=device)
dt_bias = torch.randn(HV, dtype=torch.float32, device=device)
num_slots = B + 2
ssm_states_base = torch.randn(num_slots, HV, K, V, dtype=dtype, device=device)
cache_indices = torch.arange(B, dtype=torch.int32, device=device)
query_start_loc = torch.arange(0, B * N + 1, N, dtype=torch.int32, device=device)
ssm_states_decode = ssm_states_base.clone()
outputs_decode = []
for step in range(N):
step_indices = [i * N + step for i in range(B)]
step_q = q[:, step_indices].contiguous()
step_k = k[:, step_indices].contiguous()
step_v = v[:, step_indices].contiguous()
step_a = a[step_indices].contiguous()
step_b = b[:, step_indices].contiguous()
decode_qsl = torch.arange(0, B + 1, dtype=torch.int32, device=device)
out = fused_sigmoid_gating_delta_rule_update(
A_log=A_log,
dt_bias=dt_bias,
q=step_q,
k=step_k,
v=step_v,
a=step_a,
b=step_b,
initial_state_source=ssm_states_decode,
initial_state_indices=cache_indices,
cu_seqlens=decode_qsl,
use_qk_l2norm_in_kernel=True,
softplus_beta=1.0,
softplus_threshold=20.0,
is_kda=True,
)
outputs_decode.append(out)
ssm_states_verify = ssm_states_base.clone()
intermediate_buffer = torch.zeros(
num_slots, N, HV, K, V, dtype=dtype, device=device
)
intermediate_indices = torch.arange(B, dtype=torch.int32, device=device)
out_verify = fused_sigmoid_gating_delta_rule_update(
A_log=A_log,
dt_bias=dt_bias,
q=q,
k=k,
v=v,
a=a,
b=b,
initial_state_source=ssm_states_verify,
initial_state_indices=cache_indices,
cu_seqlens=query_start_loc,
use_qk_l2norm_in_kernel=True,
softplus_beta=1.0,
softplus_threshold=20.0,
is_kda=True,
disable_state_update=True,
intermediate_states_buffer=intermediate_buffer,
intermediate_state_indices=intermediate_indices,
cache_steps=N,
retrieve_parent_token=None,
)
out_decode_list = []
for req_idx in range(B):
for step in range(N):
out_decode_list.append(outputs_decode[step][:, req_idx : req_idx + 1])
out_decode_cat = torch.cat(out_decode_list, dim=1)
max_diff = (out_verify - out_decode_cat).abs().max().item()
mean_diff = (out_verify - out_decode_cat).abs().mean().item()
print(f"\n[bf16] Output max diff: {max_diff:.6e}, mean diff: {mean_diff:.6e}")
# FP32 accumulation keeps this close to the sequential bf16 path.
assert max_diff < 1e-3, f"[bf16] Output mismatch! max diff: {max_diff}"
print("PASSED: KDA target_verify bf16 test!")
if __name__ == "__main__":
test_kda_target_verify_equivalence()
test_kda_target_verify_bf16()
test/registered/unit/spec/test_ngram_mamba_verify_update.py
+371
View File
@@ -0,0 +1,371 @@
import unittest
from unittest.mock import MagicMock, patch
import torch
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class TestNgramLastCorrectStepIndices(CustomTestCase):
def _compute_last_correct_step_indices(
self,
accept_indices: torch.Tensor,
num_correct_drafts: torch.Tensor,
draft_token_num: int,
) -> torch.Tensor:
bs = accept_indices.shape[0]
req_idx = torch.arange(bs, dtype=torch.int64, device=accept_indices.device)
accept_indices_offset = (req_idx * draft_token_num).to(accept_indices.dtype)
last_correct_step_indices = (
accept_indices[req_idx, num_correct_drafts.to(torch.int64)]
- accept_indices_offset
)
return last_correct_step_indices
def test_linear_chain_all_accepted(self):
bs, draft_token_num = 3, 5
accept_indices = torch.stack(
[
torch.arange(
i * draft_token_num,
i * draft_token_num + draft_token_num,
dtype=torch.int32,
)
for i in range(bs)
]
)
num_correct_drafts = torch.tensor([4, 4, 4], dtype=torch.int32)
result = self._compute_last_correct_step_indices(
accept_indices, num_correct_drafts, draft_token_num
)
expected = torch.tensor([4, 4, 4], dtype=torch.int32)
self.assertTrue(torch.equal(result, expected))
def test_linear_chain_partial_accept(self):
bs, draft_token_num = 3, 5
accept_indices = torch.tensor(
[
[0, 1, 2, -1, -1],
[5, -1, -1, -1, -1],
[10, 11, 12, 13, 14],
],
dtype=torch.int32,
)
num_correct_drafts = torch.tensor([2, 0, 4], dtype=torch.int32)
result = self._compute_last_correct_step_indices(
accept_indices, num_correct_drafts, draft_token_num
)
expected = torch.tensor([2, 0, 4], dtype=torch.int32)
self.assertTrue(torch.equal(result, expected))
def test_tree_structure_non_sequential(self):
bs, draft_token_num = 2, 6
accept_indices = torch.tensor(
[
[0, 2, 5, -1, -1, -1],
[6, 7, 10, -1, -1, -1],
],
dtype=torch.int32,
)
num_correct_drafts = torch.tensor([2, 2], dtype=torch.int32)
result = self._compute_last_correct_step_indices(
accept_indices, num_correct_drafts, draft_token_num
)
expected = torch.tensor([5, 4], dtype=torch.int32)
self.assertTrue(torch.equal(result, expected))
def test_single_request_zero_drafts(self):
bs, draft_token_num = 1, 4
accept_indices = torch.tensor([[0, -1, -1, -1]], dtype=torch.int32)
num_correct_drafts = torch.tensor([0], dtype=torch.int32)
result = self._compute_last_correct_step_indices(
accept_indices, num_correct_drafts, draft_token_num
)
expected = torch.tensor([0], dtype=torch.int32)
self.assertTrue(torch.equal(result, expected))
class TestNgramMambaVerifyUpdate(CustomTestCase):
def _make_mock_target_worker(self):
target_worker = MagicMock()
target_worker.model_runner.model = MagicMock()
target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify = (
MagicMock()
)
return target_worker
def test_mamba_verify_update_called_with_correct_indices(self):
from sglang.srt.speculative.spec_utils import commit_mamba_states_after_verify
target_worker = self._make_mock_target_worker()
batch = MagicMock()
batch.forward_mode.is_idle.return_value = False
batch.mamba_track_indices = None
batch.seq_lens = torch.tensor([10, 20, 30], dtype=torch.int32)
accept_lens = torch.tensor([3, 1, 5], dtype=torch.int32)
accept_index = torch.tensor(
[
[0, 1, 2, -1, -1],
[5, -1, -1, -1, -1],
[10, 11, 12, 13, 14],
],
dtype=torch.int32,
)
with patch(
"sglang.srt.speculative.spec_utils.mambaish_config",
return_value={"some": "config"},
):
commit_mamba_states_after_verify(
target_worker,
batch,
accept_lens,
accept_index,
draft_token_num=5,
)
update_call = (
target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify
)
update_call.assert_called_once()
call_kwargs = update_call.call_args[1]
self.assertTrue(
torch.equal(
call_kwargs["last_correct_step_indices"],
torch.tensor([2, 0, 4], dtype=torch.int32),
)
)
self.assertIsNone(call_kwargs["mamba_track_indices"])
self.assertIsNone(call_kwargs["mamba_steps_to_track"])
def test_mamba_verify_update_not_called_for_non_mamba_model(self):
from sglang.srt.speculative.spec_utils import commit_mamba_states_after_verify
target_worker = self._make_mock_target_worker()
batch = MagicMock()
batch.forward_mode.is_idle.return_value = False
batch.mamba_track_indices = None
accept_lens = torch.tensor([1], dtype=torch.int32)
accept_index = torch.tensor([[0, -1, -1, -1, -1]], dtype=torch.int32)
with patch(
"sglang.srt.speculative.spec_utils.mambaish_config",
return_value=None,
):
commit_mamba_states_after_verify(
target_worker,
batch,
accept_lens,
accept_index,
draft_token_num=5,
)
update_call = (
target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify
)
update_call.assert_not_called()
def test_mamba_verify_update_with_track_indices(self):
from sglang.srt.speculative.spec_utils import commit_mamba_states_after_verify
target_worker = self._make_mock_target_worker()
batch = MagicMock()
batch.forward_mode.is_idle.return_value = False
batch.mamba_track_indices = torch.tensor([100, 200], dtype=torch.int64)
# Only the first request crosses the 256-token tracking boundary.
batch.seq_lens = torch.tensor([253, 128], dtype=torch.int32)
accept_lens = torch.tensor([4, 3], dtype=torch.int32)
accept_index = torch.tensor(
[
[0, 1, 2, 3, -1],
[5, 6, 7, -1, -1],
],
dtype=torch.int32,
)
with patch(
"sglang.srt.speculative.spec_utils.mambaish_config",
return_value={"some": "config"},
), patch(
"sglang.srt.speculative.spec_utils.get_server_args",
return_value=MagicMock(mamba_track_interval=256),
):
commit_mamba_states_after_verify(
target_worker,
batch,
accept_lens,
accept_index,
draft_token_num=5,
)
update_call = (
target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify
)
update_call.assert_called_once()
call_kwargs = update_call.call_args[1]
self.assertTrue(
torch.equal(
call_kwargs["last_correct_step_indices"],
torch.tensor([3, 2], dtype=torch.int32),
)
)
self.assertTrue(
torch.equal(
call_kwargs["mamba_steps_to_track"],
torch.tensor([2, -1], dtype=torch.int32),
)
)
class TestConvWindowDedupLayout(CustomTestCase):
"""KDA stores conv_state as (K-1, channel), unlike GDN; partial-accept
commits must preserve that layout in the overlapping view.
"""
@staticmethod
def _build_fixed_view(channel_dim, win_len, draft_tokens, window_major, device):
shared_win = draft_tokens + win_len - 1
L, S = 1, 1
phys = torch.zeros(L, S, channel_dim, shared_win, device=device)
# Encoding both coordinates makes axis aliasing observable.
for c in range(channel_dim):
for w in range(shared_win):
phys[0, 0, c, w] = c * 1000 + w
if not window_major:
# GDN: view[l, s, step, d, w] = phys[l, s, d, step + w]
view = phys.as_strided(
(L, S, draft_tokens, channel_dim, win_len),
(
phys.stride(0),
phys.stride(1),
phys.stride(3),
phys.stride(2),
phys.stride(3),
),
)
else:
# KDA: view[l, s, step, w, d] = phys[l, s, d, step + w]
view = phys.as_strided(
(L, S, draft_tokens, win_len, channel_dim),
(
phys.stride(0),
phys.stride(1),
phys.stride(3),
phys.stride(3),
phys.stride(2),
),
)
return view, phys
@staticmethod
def _build_buggy_kda_view(channel_dim, win_len, draft_tokens, device):
"""Preserve the former axis swap so the regression test distinguishes
the corrected view from the broken one.
"""
conv_shape = (win_len, channel_dim)
conv_dim, win = conv_shape
shared_win = draft_tokens + win - 1
L, S = 1, 1
phys = torch.zeros(L, S, conv_dim, shared_win, device=device)
for c in range(conv_dim):
for w in range(shared_win):
phys[0, 0, c, w] = c * 1000 + w
view = phys.as_strided(
(L, S, draft_tokens, conv_dim, win),
(
phys.stride(0),
phys.stride(1),
phys.stride(3),
phys.stride(2),
phys.stride(3),
),
)
return view
def test_kda_window_major_sliding_window(self):
channel_dim, win_len, draft_tokens = 5, 3, 4
view, _ = self._build_fixed_view(
channel_dim, win_len, draft_tokens, window_major=True, device="cpu"
)
for t in range(draft_tokens):
for w in range(win_len):
for d in range(channel_dim):
got = int(view[0, 0, t, w, d].item())
self.assertEqual(
got,
d * 1000 + (t + w),
msg=f"KDA view alias at step={t} w={w} d={d}",
)
def test_kda_channel_axis_independent(self):
channel_dim, win_len, draft_tokens = 5, 3, 4
view, _ = self._build_fixed_view(
channel_dim, win_len, draft_tokens, window_major=True, device="cpu"
)
for t in range(draft_tokens):
for w in range(win_len):
for d in range(channel_dim):
self.assertEqual(int(view[0, 0, t, w, d].item()) // 1000, d)
def test_kda_window_shifts_by_one_per_step(self):
channel_dim, win_len, draft_tokens = 5, 3, 4
view, _ = self._build_fixed_view(
channel_dim, win_len, draft_tokens, window_major=True, device="cpu"
)
fixed_channel = 2
for t in range(draft_tokens - 1):
a = view[0, 0, t, :, fixed_channel].tolist()
b = view[0, 0, t + 1, :, fixed_channel].tolist()
self.assertEqual(a[1:], b[:-1])
def test_gdn_channel_major_unchanged(self):
channel_dim, win_len, draft_tokens = 5, 3, 4
view, _ = self._build_fixed_view(
channel_dim, win_len, draft_tokens, window_major=False, device="cpu"
)
for t in range(draft_tokens):
for d in range(channel_dim):
for w in range(win_len):
self.assertEqual(
int(view[0, 0, t, d, w].item()), d * 1000 + (t + w)
)
def test_partial_accept_commit_reads_correct_window(self):
channel_dim, win_len, draft_tokens = 5, 3, 4
view, _ = self._build_fixed_view(
channel_dim, win_len, draft_tokens, window_major=True, device="cpu"
)
n = 1
committed = view[0, 0, n]
for w in range(win_len):
for d in range(channel_dim):
self.assertEqual(int(committed[w, d].item()), d * 1000 + (n + w))
def test_buggy_kda_view_aliases_step_onto_channel(self):
channel_dim, win_len, draft_tokens = 5, 3, 4
buggy = self._build_buggy_kda_view(
channel_dim, win_len, draft_tokens, device="cpu"
)
self.assertEqual(buggy.shape[3], win_len)
self.assertEqual(buggy.shape[4], channel_dim)
aliased = False
for c in range(buggy.shape[3]):
if buggy[0, 0, 0, c, :].tolist() != buggy[0, 0, 1, c, :].tolist():
aliased = True
break
self.assertTrue(
aliased,
"expected the buggy KDA view to alias the draft-step axis onto the "
"channel axis",
)
if __name__ == "__main__":
unittest.main()