Barely any speedup when running Llama for inference (with Kernel Injection) [BUG]

As with other models when using DS inference with a Batch size 1. You would expect to see more inference speedup using kernel injection.
However tokens per second is very similar to vanilla Pytorch.

To Reproduce
Using Deepspeed - v0.9.2
required = HF Transformers, pytorch
This was ran on an A6000, using the latest HF transformers pytorch docker image

Script to test

import torch
import deepspeed
import time

# Load the tokenizer
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("TheBloke/stable-vicuna-13B-HF")

print("Loading original model...")
model = LlamaForCausalLM.from_pretrained("TheBloke/stable-vicuna-13B-HF", torch_dtype=torch.float16).cuda()

# Prepare the input
prompt = "this is a test prompt"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()

# Start timing for the original model
print("Generating tokens using original model...")
start_time = time.time()

tokens = model.generate(input_ids, max_new_tokens=500)

end_time = time.time()
original_model_time = end_time - start_time

# Compute and print tokens per second for original model
original_model_tps = 500 / original_model_time
print(f'Original model completed. Tokens per second: {original_model_tps}')

# Free up GPU memory
print("Clearing GPU memory...")
del model
torch.cuda.empty_cache()

# Reload the model for deepspeed
print("Reloading model for deepspeed...")
model = LlamaForCausalLM.from_pretrained("TheBloke/stable-vicuna-13B-HF", torch_dtype=torch.float16).cuda()

# Initialize deepspeed model
print("Initializing deepspeed model...")
ds_model = deepspeed.init_inference(model=model, mp_size=1, dtype=torch.float16, replace_with_kernel_inject=True).eval()

# Start timing for the deepspeed model
print("Generating tokens using deepspeed model...")
start_time = time.time()

tokens = ds_model.generate(input_ids, max_new_tokens=500)

end_time = time.time()
deepspeed_model_time = end_time - start_time

# Compute and print tokens per second for deepspeed model
deepspeed_model_tps = 500 / deepspeed_model_time
print(f'Deepspeed model completed. Tokens per second: {deepspeed_model_tps}')

Results:
Original (pytorch) model completed. Tokens per second: 49.076392832520355

Deepspeed (ds inf) model completed. Tokens per second: 50.66601410413505

Expected behavior
You should expect to see better performance from the DS inference model than we are currently getting.

ds_report output

DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variablesto where it can be found.
async_io ............... [NO] ....... [NO]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.0
 [WARNING]  using untested triton version (2.0.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
Traceback (most recent call last):
  File "/opt/conda/bin/ds_report", line 6, in <module>
    cli_main()
  File "/opt/conda/lib/python3.10/site-packages/deepspeed/env_report.py", line 121, in cli_main
    main(hide_operator_status=args.hide_operator_status, hide_errors_and_warnings=args.hide_errors_and_warnings)
  File "/opt/conda/lib/python3.10/site-packages/deepspeed/env_report.py", line 115, in main
    op_report(verbose=not hide_errors_and_warnings)
  File "/opt/conda/lib/python3.10/site-packages/deepspeed/env_report.py", line 52, in op_report
    is_compatible = OKAY if builder.is_compatible(verbose) else no
  File "/opt/conda/lib/python3.10/site-packages/deepspeed/ops/op_builder/spatial_inference.py", line 29, in is_compatible
    sys_cuda_major, _ = installed_cuda_version()
  File "/opt/conda/lib/python3.10/site-packages/deepspeed/ops/op_builder/builder.py", line 43, in installed_cuda_version
    output = subprocess.check_output([cuda_home + "/bin/nvcc", "-V"], universal_newlines=True)
  File "/opt/conda/lib/python3.10/subprocess.py", line 421, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/opt/conda/lib/python3.10/subprocess.py", line 503, in run
    with Popen(*popenargs, **kwargs) as process:
  File "/opt/conda/lib/python3.10/subprocess.py", line 971, in __init__
    self._execute_child(args, executable, preexec_fn, close_fds,
  File "/opt/conda/lib/python3.10/subprocess.py", line 1847, in _execute_child
    raise child_exception_type(errno_num, err_msg, err_filename)
NotADirectoryError: [Errno 20] Not a directory: '/opt/conda/bin/nvcc/bin/nvcc'```


**System info (please complete the following information):**
 - OS: Ubuntu 18.04
 - 1 X A6000