[BUG]something related to triton inference of TritonAttention Layer

Hi friend:
I am testing something related triton self-attention layer using triton based on this class definition: https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/ops/transformer/inference/triton/attention.py. Here is what I did for my script:

import torch
from torch import nn

from deepspeed.ops.transformer.inference.triton.attention import TritonSelfAttention
from deepspeed.ops.transformer.inference.config import DeepSpeedInferenceConfig

config = DeepSpeedInferenceConfig(hidden_size=5,
                 intermediate_size=5,
                 heads=1,
                 num_hidden_layers=1,
                 use_triton=True,
                 pre_layer_norm=False)

model = TritonSelfAttention(config)


for name, param in model.named_parameters():
    print('name', name)
    print('param', param)


data_input = torch.randn((1, 16, 32), dtype=torch.float16, device="cuda:0")
data_input = data_input.contiguous()
data_output = model(input = data_input, input_mask=None)

print(data_output.cpu().detach().numpy())

However, after trying these lines' codes, I found there is something related to cuda error (I feel it is related to something about kernel implementation? ).

!!!! kernel execution error. (m: 15, n: 16, k: 32, error: 7) 
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /home/lclcq/infer_workspace/DeepSpeed/deepspeed/ops/transformer/inference/triton/my_attention/te │
│ st.py:28 in <module>                                                                             │
│                                                                                                  │
│   25                                                                                             │
│   26 data_input = torch.randn((1, 16, 32), dtype=torch.float16, device="cuda:0")                 │
│   27 data_input = data_input.contiguous()                                                        │
│ ❱ 28 data_output = model(data_input)                                                             │
│   29 exit(0)                                                                                     │
│   30                                                                                             │
│   31                                                                                             │
│                                                                                                  │
│ /home/lclcq/miniconda3/envs/new_env2/lib/python3.8/site-packages/torch/nn/modules/module.py:1130 │
│ in _call_impl                                                                                    │
│                                                                                                  │
│   1127 │   │   # this function, and just call forward.                                           │
│   1128 │   │   if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks o  │
│   1129 │   │   │   │   or _global_forward_hooks or _global_forward_pre_hooks):                   │
│ ❱ 1130 │   │   │   return forward_call(*input, **kwargs)                                         │
│   1131 │   │   # Do not call functions when jit is used                                          │
│   1132 │   │   full_backward_hooks, non_full_backward_hooks = [], []                             │
│   1133 │   │   if self._backward_hooks or _global_backward_hooks:                                │
│                                                                                                  │
│ /home/lclcq/infer_workspace/DeepSpeed/deepspeed/ops/transformer/inference/triton/attention.py:17 │
│ 9 in forward                                                                                     │
│                                                                                                  │
│   176 │   │   │   qkv = qkv_out[0]                                                               │
│   177 │   │                                                                                      │
│   178 │   │   if use_triton_attention and (alibi is None):                                       │
│ ❱ 179 │   │   │   context_layer = compute_attention(qkv=qkv,                                     │
│   180 │   │   │   │   │   │   │   │   │   │   │     input_mask=input_mask,                       │
│   181 │   │   │   │   │   │   │   │   │   │   │     scale=self.scale,                            │
│   182 │   │   │   │   │   │   │   │   │   │   │     layer_past=layer_past,                       │
│                                                                                                  │
│ /home/lclcq/infer_workspace/DeepSpeed/deepspeed/ops/transformer/inference/triton/attention.py:22 │
│ 2 in compute_attention                                                                           │
│                                                                                                  │
│   219 │                                                                                          │
│   220 │   #assert layer_past is None, "layer_past not supported in triton yet"                   │
│   221 │   assert alibi is None, "layer_past not supported in alibi yet"                          │
│ ❱ 222 │   output = score_4d_matmul(qkv, head_size, triangular, scale)                            │
│   223 │   if triangular:                                                                         │
│   224 │   │   output = softmax(output)                                                           │
│   225 │   else:                                                                                  │
│                                                                                                  │
│ /home/lclcq/infer_workspace/DeepSpeed/deepspeed/ops/transformer/inference/triton/matmul_ext.py:2 │
│ 89 in _score_4d_matmul                                                                           │
│                                                                                                  │
│   286 │   │   # print(input)                                                                     │
│   287 │   │   # exit(0)                                                                          │
│   288 │   │                                                                                      │
│ ❱ 289 │   │   q = input[:, :, :d_model]                                                          │
│   290 │   │   # exit(0)                                                                          │
│   291 │   │                                                                                      │
│   292 │   │   k = input[:, :, d_model:d_model * 2]                                               │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: CUDA error: invalid device ordinal
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

Not quite sure whether it is related to my script or something else? Thank you for your help!

Start with the provided reproduction script and deepspeed/ops/transformer/inference/triton/attention.py, following forward through compute_attention and score_4d_matmul in matmul_ext.py. Use CUDA_LAUNCH_BLOCKING=1 while checking the reported invalid device ordinal; done means identifying whether the failure is caused by the script or the Triton implementation and verifying the resulting behavior.