Repository navigation
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
45bd3b0 to
1010d5f
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
1010d5f to
b57e115
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
b57e115 to
9e0a182
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
| on_cdna() | ||
| and find_spec("flash_attn") is not None | ||
| and (dtype == torch.float16 or dtype == torch.bfloat16) | ||
| if (on_cdna() or on_rdna()) and ( |
There was a problem hiding this comment.
you can use on_gfx1x() instead of on_rdna to just snag the 11 and 12 gfx arch's
ViT already supports FlashAttention on ROCm. Update the ROCm and TheRock Docker pins to 4a948e9c and include RDNA targets in FlashAttention CK builds. With RDNA CK support in the updated dependency, remove the Triton-only restriction from automatic ViT FlashAttention selection on RDNA. Use the same package and dtype checks as CDNA, preserving CDNA AITER priority. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Hosang Yoon <hosang.yoon@amd.com>
9e0a182 to
b062ca7
Compare
|
This pull request has merge conflicts that must be resolved before it can be |
dllehr-amd
left a comment
dllehr-amd
left a comment
There was a problem hiding this comment.
Looks good! @AndreasKaratzas @Rohan138 I'll leave it to you to decide when to suck this into our Dockerfiles. Essentially upstream flash attention (tridao) works with rdna now, so we don't need to exclude them from the build, and only need a flash attention commit bump.
|
/ci run |
|
✅ @hyoon1, CI is now available for this PR.
|
|
❌ This PR is 885 commits behind upstream |
Adopt the current ROCm Dockerfile layout and preserve the existing ViT selector change. Reapply the FlashAttention build updates in a follow-up commit. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Hosang Yoon <hosang.yoon@amd.com>
Update the ROCm 10 and ROCm 7.2 base images to FlashAttention 4a948e9c and include RDNA targets in CK builds after merging the new Dockerfile layout. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Hosang Yoon <hosang.yoon@amd.com>
|
@hyoon1 could you rebase this PR against the current HEAD? Sorry we moved the dockerfiles around a bunch last week to enable ROCm 10 |
|
/ci run |
|
/amd-ci run nightly |
|
✅ Triggered Buildkite CI #96216 for commit |
|
✅ Triggered Buildkite AMD CI #15072 for commit |
|
I'm good to merge this (pending CI) btw! It'll cause a couple minor merge conflicts but hopefully nothing crazy @dllehr-amd |
|
Hi @Rohan138, seems CI is interrupted during the base docker rebuilding. Could you please check? |
|
This pull request has merge conflicts that must be resolved before it can be |
Keep the unified ROCm base image layout and remove the retired ROCm 7.2 base Dockerfile. Preserve FlashAttention 4a948e9c, the full RDNA CK build targets, and the reviewed ViT backend selection changes. Retain upstream's build pipeline updates and the separate gfx1250 base image. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Hosang Yoon <hosang.yoon@amd.com>
|
❌ @hyoon1, your CI command failed with an unexpected error, so CI may not have started. See the workflow run for details. |
|
/ci run |
|
/amd-ci run nightly |
|
✅ Triggered Buildkite CI #100003 for commit |
|
✅ Triggered Buildkite AMD CI #15769 for commit |
Overview
Update FlashAttention to 4a948e9c in the ROCm Docker base images and include RDNA targets in CK builds, allowing the existing ViT FlashAttention path to use CK by default on RDNA3/RDNA4.
FA_BRANCHfrom0e60e394to4a948e9cindocker/Dockerfile.rocm_baseanddocker/Dockerfile.rock_base. This commit includes RDNA CK support and the CKvarlen_fwdbinding fix.PYTORCH_ROCM_ARCHlist toGPU_ARCHS.flash_attnpackage is available and the dtype is FP16 or BF16.FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUEand the private implementation module probes.Explicit backend selection and the existing CDNA AITER FlashAttention priority remain unchanged.
FLASH_ATTENTION_TRITON_AMD_ENABLEFALSETORCH_SDPAFLASH_ATTN(CK)TRUEFLASH_ATTN(Triton)FLASH_ATTN(Triton)FALSETORCH_SDPAFLASH_ATTN(Triton via FA's own fallback)Validation
Existing tests
The existing ROCm attention backend selection tests were run:
Result: 76 passed, 15 warnings.
No test files are changed in this PR.
Runtime validation
A 100-request multimodal workload using Qwen3.5-9B was validated under the following ViT attention configurations:
--mm-encoder-attn-backend TORCH_SDPAFLASH_ATTENTION_TRITON_AMD_ENABLEunset orFALSEFLASH_ATTENTION_TRITON_AMD_ENABLE=TRUERuntime validation command
vllm bench throughput \ --model /home/user/workspace/models/Qwen3.5-9B \ --trust-remote-code \ --backend vllm-chat \ --dataset-name random-mm \ --dataset-path /dev/null \ --random-mm-bucket-config '{(1024, 1024, 1): 1.0}' \ --random-input-len 1024 \ --random-output-len 128 \ --num-prompts 100 \ --num-warmups 1 \ --max-model-len 4096Apply the setting listed in the table for each configuration.
Details
ViT FlashAttention support already exists in vLLM. Previously, the Docker build excluded RDNA targets from FlashAttention CK builds, and automatic FlashAttention selection on RDNA required Triton opt-in.
Following the maintainer discussion, this revision updates the existing PR to include both the FlashAttention dependency/build changes and the corresponding selector simplification. Updating the pin alone would leave RDNA CK kernels excluded by the existing build filter.
FlashAttention handles CK/Triton implementation selection. The pinned version also falls back to Triton if importing the CK extension raises
ImportError. This fallback does not cover GPU execution errors after a successful import.Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.