[REQUEST] During inference, support passing `past_key_values` even if `input_ids.shape[-1] >= 2`

PROBLEM

I would like to use DeepSpeed together with the LLMA technique described in this Microsoft paper (cc @nyanyanya ). LLMA, which is related to speculative decoding, speeds up model inference by 3x for several of my use cases.

But I can't use DeepSpeed, because it doesn't seem to support inputting my own past_key_values if the length of input_ids is >=2. This is related to this comment by @RezaYazdaniAminabadi

CODE TO REPRODUCE

The following code shows how DeepSpeed does not produce the same output (compared to not running with DeepSpeed) when we use past_key_values in combination with input id lengths >=2.

# !pip install torch==2.0.1 transformers==4.34.0 deepspeed==0.11.1

from typing import Optional, Any
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import deepspeed


ARCHITECTURE = "gpt2"
DEVICE = "cuda"
PROMPT = "Capital city of Germany:"
MAX_NEW_TOKENS = 10
TOKENIZER = AutoTokenizer.from_pretrained(ARCHITECTURE, use_fast=True)


def cut_pkv(pkv: tuple, length: int = 0) -> tuple:
    """Return past_key_values cut to the first `length` tokens."""
    return tuple((t1[:, :, :length, :], t2[:, :, :length, :]) for (t1, t2) in pkv)

def run_WITHOUT_past_key_values(model: Any) -> None:
    input_ids = TOKENIZER(PROMPT, return_tensors="pt")["input_ids"].to(DEVICE)
    input_plus_completion_ids = input_ids

    with torch.inference_mode():
        for _ in range(MAX_NEW_TOKENS):
            output = model(input_plus_completion_ids)
            predicted_id = torch.argmax(output.logits[0][-1])
            input_plus_completion_ids = torch.cat((input_plus_completion_ids, predicted_id.view(1, 1)), dim=1)

    print(TOKENIZER.decode(input_plus_completion_ids[0]))  
    
    
def run_WITH_past_key_values(model: Any) -> None:
    input_ids = TOKENIZER(PROMPT, return_tensors="pt")["input_ids"].to(DEVICE)
    input_plus_completion_ids = input_ids
    pkv = None

    with torch.inference_mode():
        for _ in range(MAX_NEW_TOKENS):
            if pkv is None:
                output = model(input_plus_completion_ids)
            else:
                # IMPORTANT: we remove all but the last 2 input ids, and pass in all past_key_values except those of the last 2 input ids
                pkv = cut_pkv(pkv, length=input_plus_completion_ids[0].shape[-1] - 2)
                output = model(input_plus_completion_ids[0:1, -2:], past_key_values=pkv, use_cache=True)


            predicted_id = torch.argmax(output.logits[0][-1])
            input_plus_completion_ids = torch.cat((input_plus_completion_ids, predicted_id.view(1, 1)), dim=1)

            pkv = output.past_key_values
            
    print(TOKENIZER.decode(input_plus_completion_ids[0]))  


# running without DeepSpeed works fine in both cases
model = AutoModelForCausalLM.from_pretrained(ARCHITECTURE).eval().to(DEVICE)
run_WITHOUT_past_key_values(model) # correctly prints out: `Capital city of Germany:\n\nThe city of Berlin is the capital of`
print()
print()
run_WITH_past_key_values(model) # correctly prints out: `Capital city of Germany:\n\nThe city of Berlin is the capital of`

# running with DeepSpeed produces an erroneous output in the second case
ds_engine = deepspeed.init_inference(model, dtype=torch.float32, replace_with_kernel_inject=True)
model = ds_engine.module
run_WITHOUT_past_key_values(model) # correctly prints out: `Capital city of Germany:\n\nThe city of Berlin is the capital of`
print()
print()
run_WITH_past_key_values(model) # WRONGLY prints out: `Capital city of Germany:\n\n:::::::::`

SOLUTION

I would like to get the same results when running a model forward pass, independently on whether I run with or without DeepSpeed.

ADDITIONAL NOTES

This feature request would likely also fix/resolve this bug and this issue/feature request mentioned by @sakoush

I'd be happy to discuss, debug and/or help implementing this.

Start by running the provided GPT-2 reproduction with the pinned torch, transformers, and DeepSpeed versions, comparing ordinary inference with the supplied past_key_values path. Trace the DeepSpeed inference forward path used by kernel injection for multi-token input_ids and verify completion when outputs match non-DeepSpeed inference, including the related cases linked in the issue.