ColossalAI/tests/test_infer_ops/triton
Yuanheng Zhao 07b5283b6a [kernel] Add triton kernel for context attention (FAv2) without padding (#5192)
* add context attn unpadded triton kernel

* test compatibility

* kv cache copy (testing)

* fix k/v cache copy

* fix kv cache copy and test

* fix boundary of block ptrs

* add support for GQA/MQA and testing

* fix import statement

---------

Co-authored-by: Round Heng <yuanhengzhao@Rounds-MacBook-Pro.local>
2024-01-11 13:39:56 +00:00
..
kernel_utils.py [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00
test_context_attn_unpad.py [kernel] Add triton kernel for context attention (FAv2) without padding (#5192) 2024-01-11 13:39:56 +00:00
test_layernorm_triton.py [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00
test_llama_act_combine.py [moe] merge moe into main (#4978) 2023-11-02 02:21:24 +00:00
test_softmax.py [misc] update pre-commit and run all files (#4752) 2023-09-19 14:20:26 +08:00