Skip to content

Commit ba1aa1e

Browse files
lonelygshguanshihui]
lonelygsh
and
guanshihui]
authored
[BugFix][Speculative Decoding] Correct index calculation in speculate decoding operators (#7121)
- Fix accept_idx calculation in spec_set_value_by_stop_seqs - Fix condition check from < to <= for token matching - Fix accept_tokens indexing logic - Remove unnecessary -1 in current_step comparison for max_think_len Co-authored-by: guanshihui] <guanshihui@baidu.com>
1 parent 7a2e330 commit ba1aa1e

3 files changed

Lines changed: 34 additions & 35 deletions

File tree

‎custom_ops/gpu_ops/speculate_decoding/speculate_limit_thinking_content_length.cu‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -98,10 +98,7 @@ __global__ void speculate_limit_thinking_content_length_kernel(
9898
if (max_think_len > 0) {
9999
// A) 超长触发：到达 max_think_len 时开始注入（从本 token 起输出
100100
// inject_token_ids[0]）
101-
if (status == 0 &&
102-
(current_step - 1) ==
103-
max_think_len) { // current_step - 1 是因为 speculate_verify 里
104-
// step_idx + 1 了
101+
if (status == 0 && current_step == max_think_len) {
105102
status = (inject_len > 0) ? 1 : done_status;
106103
}
107104
} else if (max_think_len == 0) {

‎custom_ops/gpu_ops/speculate_decoding/speculate_set_stop_value_multi_seqs.cu‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ __global__ void spec_set_value_by_stop_seqs(bool *stop_flags,
5858
bool is_end = false;
5959
// 遍历起始位置
6060
for (; accept_idx <= accept_num - 1 && !is_end; accept_idx++) {
61-
if (step_idx_now + accept_idx + 1 < stop_seq_len) {
61+
if (step_idx_now - accept_num + accept_idx + 1 < stop_seq_len) {
6262
#ifdef DEBUG_SPEC_STOP_SEQS
6363
printf("num %d < stop_seq_len %d\n",
6464
step_idx_now - accept_num + accept_idx + 1,
@@ -71,7 +71,7 @@ __global__ void spec_set_value_by_stop_seqs(bool *stop_flags,
7171
int64_t cur_token_idx = -1;
7272

7373
// 通过当前值判断 token 是在 pre_ids 还是 accept_token 里
74-
if (stop_seq_len - 1 - i < accept_idx) {
74+
if (stop_seq_len - 1 - i <= accept_idx) {
7575
#ifdef DEBUG_SPEC_STOP_SEQS
7676
printf(
7777
"AcceptTokens bid:%d. tid:%d, accept_idx:%d, "
@@ -83,7 +83,7 @@ __global__ void spec_set_value_by_stop_seqs(bool *stop_flags,
8383
accept_idx - (stop_seq_len - 1 - i) - 1);
8484
#endif
8585
cur_token_idx =
86-
accept_tokens_now[accept_idx - (stop_seq_len - 1 - i) - 1];
86+
accept_tokens_now[accept_idx - (stop_seq_len - 1 - i)];
8787
} else {
8888
#ifdef DEBUG_SPEC_STOP_SEQS
8989
printf(
@@ -98,7 +98,7 @@ __global__ void spec_set_value_by_stop_seqs(bool *stop_flags,
9898
(stop_seq_len - 1 - i));
9999
#endif
100100
int pre_ids_idx =
101-
step_idx_now + accept_idx - (stop_seq_len - 1 - i);
101+
step_idx_now - accept_num + accept_idx - (stop_seq_len - 1 - i);
102102
// EC3
103103
// 特殊拼接会导致input_ids最后一位无特殊token，即pre_ids[0]可能为23,
104104
// 导致异常结束

‎tests/operators/test_speculate_set_stop_value_multi_seqs.py‎

Lines changed: 29 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -175,17 +175,19 @@ def reference_spec_set_stop_value_multi_seqs(inputs: Dict[str, Any]) -> Dict[str
175175
accept_idx = 0
176176
is_end = False
177177
while accept_idx <= an - 1 and not is_end:
178-
if step_idx_now + accept_idx + 1 < stop_seq_len:
178+
if step_idx_now - an + accept_idx + 1 < stop_seq_len:
179179
accept_idx += 1
180180
continue
181181

182182
# Check one stop_seq match
183183
for i in range(stop_seq_len - 1, -1, -1):
184184
cur_token_idx = -1
185-
if stop_seq_len - 1 - i < accept_idx:
186-
cur_token_idx = accept_tokens_now[accept_idx - (stop_seq_len - 1 - i) - 1]
185+
# 注意:新版本kernel改成了 <=,并且去掉了 -1
186+
if stop_seq_len - 1 - i <= accept_idx:
187+
cur_token_idx = accept_tokens_now[accept_idx - (stop_seq_len - 1 - i)]
187188
else:
188-
pre_ids_idx = step_idx_now + accept_idx - (stop_seq_len - 1 - i)
189+
# 新版本:step_idx已经包含accept_num,所以要减去
190+
pre_ids_idx = step_idx_now - an + accept_idx - (stop_seq_len - 1 - i)
189191
if pre_ids_idx <= 0:
190192
break
191193
cur_token_idx = pre_ids_now[pre_ids_idx]
@@ -290,22 +292,22 @@ def test_match_spanning_pre_ids_and_accept(self):
290292
inputs["prompt_lens"][:] = 0
291293
inputs["step_idx"][:] = 6
292294
inputs["accept_num"][:] = 3
293-
# Kernel matching at accept_idx=2 (3rd token, 0-indexed):
294-
# i=2(last): stop_seq_len-1-i=0 < accept_idx(2) -> accept_tokens[2-0-1]=accept_tokens[1]
295-
# i=1: stop_seq_len-1-i=1 < accept_idx(2) -> accept_tokens[2-1-1]=accept_tokens[0]
296-
# i=0: stop_seq_len-1-i=2 >= accept_idx(2) -> pre_ids[step_idx+2-(3-1-0)]=pre_ids[6]
297-
# So stop_seq should be [pre_ids[6], accept_tokens[0], accept_tokens[1]]
298-
inputs["token_ids_all"][0, 6] = 99
295+
# stop_seq spans pre_ids and accept_tokens
296+
# For accept_idx=1: step_idx_now - accept_num + 1 + 1 = 6-3+1+1 = 5 >= stop_seq_len=3, so we check
297+
# i=2: stop_seq_len-1-i=0 <= accept_idx(1) -> accept_tokens[1-0] = accept_tokens[1] = 22
298+
# i=1: stop_seq_len-1-i=1 <= accept_idx(1) -> accept_tokens[1-1] = accept_tokens[0] = 11
299+
# i=0: stop_seq_len-1-i=2 > accept_idx(1) -> pre_ids_idx = 6-3+1-(3-1-0) = 4-2 = 2 -> pre_ids[2] = 99
300+
inputs["token_ids_all"][0, 2] = 99
299301
inputs["accept_tokens"][0, :3] = [11, 22, 33]
300302
inputs["stop_seqs"][0, 0, :3] = [99, 11, 22]
301303
inputs["stop_seqs_len"][0, 0] = 3
302304
inputs["stop_flags"][:] = False
303305
inputs["min_tokens"][:] = 0
304306
outputs = self._run_and_get(inputs)
305307
self._check_all_outputs(inputs, outputs)
306-
# Match at accept_idx=2, loop increments to 3
307-
self.assertEqual(outputs["accept_num"][0], 3)
308-
self.assertEqual(outputs["accept_tokens"][0, 2], -1)
308+
# Match at accept_idx=1, loop increments to 2
309+
self.assertEqual(outputs["accept_num"][0], 2)
310+
self.assertEqual(outputs["accept_tokens"][0, 1], inputs["end_ids"][0])
309311

310312
def test_match_in_pre_ids_only(self):
311313
"""Stop seq found entirely within token_ids_all (pre_ids), matching at accept_idx=0."""
@@ -314,29 +316,29 @@ def test_match_in_pre_ids_only(self):
314316
accept_tokens_len=5,
315317
max_model_len=32,
316318
stop_seqs_bs=1,
317-
stop_seqs_max_len=3,
319+
stop_seqs_max_len=4, # 需要4个元素
318320
seed=30,
319321
)
320322
inputs["prompt_lens"][:] = 0
321323
inputs["step_idx"][:] = 8
322324
inputs["accept_num"][:] = 3
323-
# pre_ids at step_idx positions: token_ids_all[0, 6]=50, [0,7]=60, [0,8]=70
324-
# stop_seq = [50, 60, 70], all 3 tokens are in pre_ids
325-
# For accept_idx=0: step_idx_now + 0 + 1 = 9 >= stop_seq_len=3, so we check
326-
# i=2: pre_ids_idx = 8+0-(3-1-2) = 8 -> pre_ids_now[8] = 70
327-
# i=1: pre_ids_idx = 8+0-(3-1-1) = 7 -> pre_ids_now[7] = 60
328-
# i=0: pre_ids_idx = 8+0-(3-1-0) = 6 -> pre_ids_now[6] = 50
329-
inputs["token_ids_all"][0, 6] = 50
330-
inputs["token_ids_all"][0, 7] = 60
331-
inputs["token_ids_all"][0, 8] = 70
332-
inputs["accept_tokens"][0, :3] = [1, 2, 3]
333-
inputs["stop_seqs"][0, 0, :3] = [50, 60, 70]
334-
inputs["stop_seqs_len"][0, 0] = 3
325+
# stop_seq partially in pre_ids, partially in accept_tokens
326+
# For accept_idx=1: step_idx_now - accept_num + 1 + 1 = 8-3+1+1 = 7 >= stop_seq_len=4, so we check
327+
# i=3: stop_seq_len-1-i=0 <= accept_idx(1) -> accept_tokens[1-0] = accept_tokens[1] = 22
328+
# i=2: stop_seq_len-1-i=1 <= accept_idx(1) -> accept_tokens[1-1] = accept_tokens[0] = 11
329+
# i=1: stop_seq_len-1-i=2 > accept_idx(1) -> pre_ids_idx = 8-3+1-(4-1-1) = 6-2 = 4 -> pre_ids[4] = 60
330+
# i=0: stop_seq_len-1-i=3 > accept_idx(1) -> pre_ids_idx = 8-3+1-(4-1-0) = 6-3 = 3 -> pre_ids[3] = 50
331+
inputs["token_ids_all"][0, 3] = 50
332+
inputs["token_ids_all"][0, 4] = 60
333+
inputs["accept_tokens"][0, :3] = [11, 22, 3]
334+
inputs["stop_seqs"][0, 0, :4] = [50, 60, 11, 22]
335+
inputs["stop_seqs_len"][0, 0] = 4
335336
inputs["stop_flags"][:] = False
336337
inputs["min_tokens"][:] = 0
337338
outputs = self._run_and_get(inputs)
338339
self._check_all_outputs(inputs, outputs)
339-
self.assertEqual(outputs["accept_num"][0], 1)
340+
# Match at accept_idx=1, loop increments to 2
341+
self.assertEqual(outputs["accept_num"][0], 2)
340342

341343
def test_already_stopped(self):
342344
"""Kernel skips sequences with stop_flags=True."""

0 commit comments

Comments
 (0)