[KDA] Add target_verify support for speculative decoding (#26888)
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
co-authored by
yuyanqi
Claude Opus 4.6
Xinyuan Tong
yuyanqi
Claude Opus 4.6
Xinyuan Tong
parent
7c4b22fae5
commit
a678a42033
@@ -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!")
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user