Feature or Bug: Is OGA support the feature all DeviceSpan<int32_t>& next_tokens are model.pad_token_id?

[buyuliushen]

Describe the bug

Is OGA support the feature all DeviceSpan<int32_t>& next_tokens are model.pad_token_id?

If the answer is Yes, there maybe bug in the following code. if the answer is No, do you have plan to support this feature? or how do you deal with this scene?

If all DeviceSpan<int32_t>& next_tokens are model.pad_token_id:
in the following code, after the end of while loop, token_index will be -1, and input_sequence_lengths[b] will be 0;

onnxruntime-genai/src/models/logits.cpp

Lines 96 to 103 in b247a93

	size_t token_index = new_kv_length;
	while (token_index-- > 0) {
	auto next_token = const_cast<DeviceSpan<int32_t>&>(next_tokens).CpuSpan()[b * new_kv_length + token_index];
	if (next_token != model_.config_->model.pad_token_id)
	break;
	}
	input_sequence_lengths[b] = static_cast<int>(token_index + 1);
	}

then go to the following code, token_index = -1; then in L55, the first parameter of logits_raw.subspan(),
(vocab_index * seq_length + token_index * vocab_size) * element_size maybe = (0 * seq_length + (-1 ) * vocab_size) * element_size = -1 * vocab_size * element_size, it's a negative value.

onnxruntime-genai/src/models/logits.cpp

Lines 52 to 58 in b247a93

	size_t token_index = input_sequence_lengths[batch_index] - 1;
	for (int beam_index = 0; beam_index < num_beams; beam_index++) {
	auto target = logits_last_tokens.subspan(vocab_index * element_size, vocab_size * element_size);
	auto source = logits_raw.subspan((vocab_index * seq_length + token_index * vocab_size) * element_size, vocab_size * element_size);
	target.CopyFrom(source);
	vocab_index += vocab_size;
	}

in the subspan(), the parameter datatype is size_t, it's unsigned, so if the first parameter is a negative value, it maybe very large as to size_t, then L417 will report error and lead to crash.

using size_type = size_t;

[RyanUnderhill]

@aciddelgado You probably know the most about this new continuous decoding logic, we should support this edge case. I can add a unit test but what do you think the best behavior is when it happens?

[buyuliushen]

Hi @RyanUnderhill @aciddelgado ， any update about this?