Repository navigation
flash-attn2: fix forward epilogue reduction for sycl-tla 0.9.2 - #1065
Conversation
The row-max reduction passed ReduceFragARow directly to cute::transform, which no longer resolves against sycl-tla 0.9.2. Access the underlying tensor explicitly and write through the Base subobject. No functional change.
|
Hi @danieldk . This PR can be reviewed now. Thanks! |
|
/kernel-bot build flash-attn2 |
|
Build request processed. Command: Dispatched (1):
|
sycl-tla 0.9.2 changed cute::reduce() to return a SubgroupTensor, so the epilogue reduction had to reach through to the underlying cute::Tensor. That form only compiles against 0.9.2: on 0.8 and 0.9.1 cute::reduce() still returns a plain cute::Tensor, which has no nested Base type, and the XPU build fails with "no type named 'Base'". Detect the subgroup wrapper instead of assuming it. as_plain_tensor() casts to the Base tensor when the type has one and is the identity otherwise, so the same source builds against every sycl-tla the build matrix uses.
|
/kernel-bot build flash-attn2 |
|
Build request processed. Command: Dispatched (1):
|
|
/kernel-bot merge-and-upload flash-attn2 |
|
Build request processed. Command: Merge result: PR is already merged. Continuing with build/upload. Dispatched (1):
Hub uploads: |
2 participants
Depends on the build-side changes in
huggingface/kernels#734.
Only a build fix: one reduction in the forward epilogue needed explicit tensor
access to compile against sycl-tla 0.9.2.
Validation
Built (
torch213-cxx11-xpu20260-x86_64-linux) and checked numerics on BMGhardware and the CRI TBX simulator. Each case runs only the kernel on XPU; data
generation, the fp32 reference and error statistics stay on CPU.
abs error 5.4e-04 (fp16) / 4.0e-03 (bf16).