From a678a42033e79a9bc3204075f2178f88a7d7f8c9 Mon Sep 17 00:00:00 2001 From: yyqjwyy <109524944+yyq0210@users.noreply.github.com> Date: Sat, 25 Jul 2026 19:52:44 +0800 Subject: [PATCH] [KDA] Add target_verify support for speculative decoding (#26888) Co-authored-by: yuyanqi Co-authored-by: Claude Opus 4.6 Co-authored-by: Xinyuan Tong --- .../layers/attention/linear/kda_backend.py | 5 + python/sglang/srt/mem_cache/memory_pool.py | 32 ++ python/sglang/srt/models/kimi_linear.py | 10 +- test/manual/test_kda_spec_integration.py | 71 ++++ test/manual/test_kda_target_verify.py | 216 ++++++++++ .../spec/test_ngram_mamba_verify_update.py | 371 ++++++++++++++++++ 6 files changed, 701 insertions(+), 4 deletions(-) create mode 100644 test/manual/test_kda_spec_integration.py create mode 100644 test/manual/test_kda_target_verify.py create mode 100644 test/registered/unit/spec/test_ngram_mamba_verify_update.py diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index 4522804d8..9c91d0584 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -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: diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 7ff2e4661..d10115143 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -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( diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 291cb208e..7ddc4bfc1 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -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] diff --git a/test/manual/test_kda_spec_integration.py b/test/manual/test_kda_spec_integration.py new file mode 100644 index 000000000..fcc96a5b4 --- /dev/null +++ b/test/manual/test_kda_spec_integration.py @@ -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!") diff --git a/test/manual/test_kda_target_verify.py b/test/manual/test_kda_target_verify.py new file mode 100644 index 000000000..a25ad5452 --- /dev/null +++ b/test/manual/test_kda_target_verify.py @@ -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() diff --git a/test/registered/unit/spec/test_ngram_mamba_verify_update.py b/test/registered/unit/spec/test_ngram_mamba_verify_update.py new file mode 100644 index 000000000..57db43e7f --- /dev/null +++ b/test/registered/unit/spec/test_ngram_mamba_verify_update.py @@ -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()