Reduce the memory usage of logits from O(context_length) to O(1)

Reduce the memory usage of logits from O(context_length) to O(1) #18038