-
Notifications
You must be signed in to change notification settings - Fork 58
[Attention] Handle fully masked reference rows #2448
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from 3 commits
Commits
Show all changes
5 commits
Select commit
Hold shift + click to select a range
8fd1528
[Attention] Handle fully masked reference rows
umangyadav 174c5e2
[Attention] Strengthen masked-row regression checks
umangyadav c0e1fbd
[Attention] Add fully masked E2E coverage
umangyadav 31ec9b3
[Rock] Guard inactive blockwise reduction threads (#2447)
umangyadav 893f0f5
add nrlargerthanthread test
umangyadav File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
19 changes: 19 additions & 0 deletions
19
mlir/test/fusion/pr-e2e/attention/rock-attention-fully-masked.mlir
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,19 @@ | ||
| // RUN: rocmlir-gen --arch %arch --operation attention -current_seq_len=2 -sliding_window_size=1 --causal -return_lse -seq_len_q 1 -seq_len_k 64 -head_dim_qk 32 -head_dim_v 32 -t f32 -rand 1 -rand_type float -pv \ | ||
| // RUN: | rocmlir-driver --host-pipeline=highlevel \ | ||
| // RUN: | rocmlir-driver -c \ | ||
| // RUN: | mlir-runner -O2 --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext,%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext,%linalg_test_lib_dir/libmlir_float16_utils%shlibext --entry-point-result=void \ | ||
| // RUN: | FileCheck %s --check-prefix=DIRECT | ||
| // RUN: rocmlir-gen --arch %arch --operation attention -current_seq_len=2 -sliding_window_size=1 --causal -return_lse -split_kv=8 -seq_len_q 1 -seq_len_k 64 -head_dim_qk 32 -head_dim_v 32 -t f32 -rand 1 -rand_type float -pv \ | ||
| // RUN: | rocmlir-driver --host-pipeline=highlevel \ | ||
| // RUN: | rocmlir-driver -c \ | ||
| // RUN: | mlir-runner -O2 --shared-libs=%linalg_test_lib_dir/libmlir_rocm_runtime%shlibext,%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext,%linalg_test_lib_dir/libmlir_float16_utils%shlibext --entry-point-result=void \ | ||
| // RUN: | FileCheck %s --check-prefix=SPLITKV | ||
|
|
||
| // A causal sliding window can leave a query with no eligible keys. Both the | ||
| // GPU kernel and the CPU reference define that row's contribution as zero. | ||
| // The direct path also returns the fully masked row's -inf LSE; the split-KV | ||
| // path exercises host recombination of fully masked partial results. | ||
|
|
||
| // DIRECT: [1 1 1] | ||
| // DIRECT-NEXT: [1 1 1] | ||
| // SPLITKV: [1 1 1] |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,37 @@ | ||
| // REQUIRES: rocm-runner | ||
| // RUN: rocmlir-driver --host-pipeline=runner %s \ | ||
| // RUN: | mlir-runner -O2 --shared-libs=%conv_validation_wrapper_library_dir/libconv-validation-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void \ | ||
| // RUN: | FileCheck %s | ||
|
|
||
| // Identical all-zero tensors have zero absolute error and zero scale. Verify | ||
| // that normalized RMS treats this exact match as zero instead of computing 0/0. | ||
|
|
||
| // CHECK: [1 1 1] | ||
|
|
||
| module { | ||
| func.func @main() { | ||
| %gpu = memref.alloc() : memref<4xf32> | ||
| %reference = memref.alloc() : memref<4xf32> | ||
| %zero = arith.constant 0.0 : f32 | ||
| linalg.fill ins(%zero : f32) outs(%gpu : memref<4xf32>) | ||
| linalg.fill ins(%zero : f32) outs(%reference : memref<4xf32>) | ||
|
|
||
| %gpuDynamic = memref.cast %gpu : memref<4xf32> to memref<?xf32> | ||
| %referenceDynamic = memref.cast %reference : memref<4xf32> to memref<?xf32> | ||
| %threshold = arith.constant 0.0 : f32 | ||
| %printDebug = arith.constant 0 : i8 | ||
| %isFP32 = arith.constant true | ||
| %useAbsDiffGate = arith.constant false | ||
| call @mcpuVerifyFloat( | ||
| %gpuDynamic, %referenceDynamic, %threshold, %threshold, %threshold, | ||
| %printDebug, %isFP32, %useAbsDiffGate) | ||
| : (memref<?xf32>, memref<?xf32>, f32, f32, f32, i8, i1, i1) -> () | ||
|
|
||
| memref.dealloc %gpu : memref<4xf32> | ||
| memref.dealloc %reference : memref<4xf32> | ||
| return | ||
| } | ||
|
|
||
| func.func private @mcpuVerifyFloat(memref<?xf32>, memref<?xf32>, f32, f32, | ||
| f32, i8, i1, i1) | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This new test actually executes, but it has no feature gate.
mlir/test/CMakeLists.txt:124-137only addsmlir-runner,mlir_runner_utils, andconv-validation-wrapperstoROCMLIR_TEST_DEPENDSwhenMLIR_ENABLE_ROCM_RUNNERis set, andmlir/test/rocmlir-driver/lit.local.cfghas no directory-level guard (unlikemlir/test/e2e/lit.local.cfg, which setsconfig.unsupported). So in a build configured without the ROCm runner this test lands in the always-oncheck-rocmlirsuite and fails on a missing tool / missing shared library rather than being skipped. Please add// REQUIRES: rocm-runneras the first line of the file (the feature is registered atmlir/test/lit.site.cfg.py.in:44-46). Note this is a pure host-side verifier test, sorocm-runneris the correct gate — no GPU or e2e gate is needed. Closest checklist bullet is the Minor "Lit test missing// RUN:line,-verify-diagnostics, or FileCheck prefix coverage", raised to Major here because the consequence is a broken test suite in a supported build configuration.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Resolved -- addressed in this revision.