register custom op for flash attn and use from torch.ops #7536
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge). To run full CI, you can do one of these:
🚀 |
| from vllm_flash_attn import flash_attn_with_kvcache as _flash_attn_with_kvcache | ||
|
|
||
|
|
||
| @torch.library.custom_op("vllm::flash_attn_varlen_func", mutates_args=[]) |
There was a problem hiding this comment.
confirmed with @WoosukKwon , these two functions do not mutate the input.
tests/kernels/test_flash_attn.py
Outdated
| cache_seqlens=kv_lens_tensor, | ||
| softcap=soft_cap if soft_cap is not None else 0, | ||
| ), | ||
| test_utils=("test_faketensor", )) |
There was a problem hiding this comment.
we can add test_schema later, after solving the OOM issue.
currently, I will get OOM when I test the schema, even though I'm using H100 80GB.
There was a problem hiding this comment.
I think part of testing the schema involves copying all the inputs and doubling checking that they are (or aren't) mutated in agreement with the op schema. I don't know if that's what is causing the OOMs here tho.
…t#7536) Signed-off-by: Alvant <[email protected]>
No description provided.