wehub-resource-sync
33 lines
865 B
Markdown
33 lines
865 B
Markdown
# a fast implementation of linear attention
|
|
|
|
## 64x64, fp16
|
|
|
|
```bash
|
|
# validate correctness
|
|
## fp16 vs fp32
|
|
python -m develop_triton_litemla attn_type=LiteMLA test_correctness=True
|
|
## triton fp16 vs fp32
|
|
python -m develop_triton_litemla attn_type=TritonLiteMLA test_correctness=True
|
|
|
|
# test performance
|
|
## fp16, forward
|
|
python -m develop_triton_litemla attn_type=LiteMLA
|
|
each step takes 10.81 ms
|
|
max memory allocated: 2.2984 GB
|
|
|
|
## triton fp16, forward
|
|
python -m develop_triton_litemla attn_type=TritonLiteMLA
|
|
each step takes 4.70 ms
|
|
max memory allocated: 1.6480 GB
|
|
|
|
## fp16, backward
|
|
python -m develop_triton_litemla attn_type=LiteMLA backward=True
|
|
each step takes 35.34 ms
|
|
max memory allocated: 3.4412 GB
|
|
|
|
## triton fp16, backward
|
|
python -m develop_triton_litemla attn_type=TritonLiteMLA backward=True
|
|
each step takes 14.25 ms
|
|
max memory allocated: 2.4704 GB
|
|
```
|