Add basic FP8 KV cache support #2603
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
danieldk
force-pushed
the
feature/fp8-kv-cache
branch
2 times, most recently
from
2628268 to
37df2ff
Compare
| key_cache.view(-1, shape[-2], shape[-1])[slots] = key | ||
| value_cache.view(-1, shape[-2], shape[-1])[slots] = value | ||
| else: | ||
| reshape_and_cache(key, value, key_cache, value_cache, slots, "auto", 1.0) |
There was a problem hiding this comment.
"auto", 1.0 ? What are those flags ? They didn't seem to be used before, aren't they defaulted in paged ?
There was a problem hiding this comment.
Specifying them here breaks IPEX no ?
There was a problem hiding this comment.
Ah, yes, the others have these arguments, but not IPEX, so I reverted this part of the PR (the KV cache now uses the existing reshape_and_cache wrappers.
danieldk
force-pushed
the
feature/fp8-kv-cache
branch
from
37df2ff to
4cc5405
Compare
This change adds rudimentary FP8 KV cache support. The support is enabled by passing `--kv-cache-dtype fp8_e5m2` to the launcher. Doing so uses this type for the KV cache. However support is still limited: * Only the `fp8_e5m2` type is supported. * The KV cache layout is the same as `float16`/`bfloat16` (HND). * The FP8 KV cache is only supported for FlashInfer. * Loading of scales is not yet supported.
danieldk
force-pushed
the
feature/fp8-kv-cache
branch
from
4cc5405 to
ed5c2fb
Compare
| kv_cache.store( | ||
| key=kv[:, :, 0].contiguous(), value=kv[:, :, 1].contiguous(), slots=slots |
There was a problem hiding this comment.
not important, but its strange we need the .contiguous() calls here
There was a problem hiding this comment.
Yeah, I wasn't sure if it was needed, these seems to be key/value striding in other places that is non-contiguous, but I also didn't want to touch it.
What does this PR do?
This change adds rudimentary FP8 KV cache support. The support is enabled by passing
--kv-cache-dtype fp8_e5m2to the launcher. Doing so uses this type for the KV cache. However support is still limited:fp8_e5m2type is supported.float16/bfloat16(HND).This PR is intentionally small to keep things reviewable. I'll follow it up with PRs that add more functionality.
Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.