feat: ROCm flash-attn varlen, triton layer norm, and AMD Dockerfile#860
feat: ROCm flash-attn varlen, triton layer norm, and AMD Dockerfile#860Abdennacer-Badaoui wants to merge 10 commits into
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. |
Abdennacer-Badaoui
force-pushed
the
fa-varlen
branch
from
0caea18 to
3c75990
Compare
alvarobartt
left a comment
alvarobartt
left a comment
There was a problem hiding this comment.
Hey @Abdennacer-Badaoui, thanks a lot, left some comments / thoughts, feel free to apply or discard where applicable 🤗
| if is_rocm(): | ||
| try: | ||
| from kernels import get_kernel as _get_kernel | ||
| _triton_layer_norm = _get_kernel("kernels-community/triton-layer-norm") |
There was a problem hiding this comment.
Can we specify the revision / version here?
| if is_rocm(): | ||
| try: | ||
| from kernels import get_kernel as _get_kernel | ||
| _triton_layer_norm = _get_kernel("kernels-community/triton-layer-norm") |
There was a problem hiding this comment.
Same here with the revision / version
| -v $volume:/data \ | ||
| --pull always \ | ||
| ghcr.io/huggingface/text-embeddings-inference:rocm-latest \ | ||
| --model-id $model --dtype bfloat16 |
There was a problem hiding this comment.
Since it just supports at most a single accelerator, could you instead specify the --device /dev/dri/renderD128 instead of --device /dev/dri or rather using the environment variable ROCR_VISIBLE_DEVICES (as per https://rocm.docs.amd.com/en/latest/conceptual/gpu-isolation.html#rocr-visible-devices)?
There was a problem hiding this comment.
Yes I agree. Let me change this
| curl http://localhost:8080/embed \ | ||
| -X POST \ | ||
| -H 'Content-Type: application/json' \ | ||
| -d '{"inputs": "What is Deep Learning?"}' |
There was a problem hiding this comment.
Might be better to use the OpenAI-compatible API at /v1/embeddings instead
| "imageNamePrefix": "rocm-", | ||
| "runOn": "always", | ||
| "sccache": true, | ||
| "extraBuildArgs": "GPU_ARCHS=gfx942", |
There was a problem hiding this comment.
Should we also build with GPU_ARCHS=gfx950, or does this image works for those too?
There was a problem hiding this comment.
yes, gfx950 already works, no separate build needed
| "dockerfile": "Dockerfile-intel" | ||
| }, | ||
| { | ||
| "name": "rocm-mi300-mi325", |
There was a problem hiding this comment.
Maybe this fits better here?
| "name": "rocm-mi300-mi325", | |
| "name": "rocm-gfx942", |
There was a problem hiding this comment.
I went with rocm-gfx942-gfx950 rather than just rocm-gfx942, since the prebuilt flash-attn wheel (whl/gfx942-gfx950/...) covers both
3 participants
Adds flash attention varlen, triton-based layer norm/rms norm via
kernels-community/triton-layer-norm, and a dedicatedDockerfile-amdwith CI matrix entry for MI300/MI325 GPUs.