dongluw

If i understand correctly, the kernel should work if the provided `ids` are output tokens instead of input tokens at step t. However I see it's called with [input tokens...