[PR #269] Fix singleton strides in CUDA flash decode #274

New issue

Open

opened 2026-10-26 22:22:20 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:22:20 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/269
Author: @curealb
Created: 10/22/2026
Status: 🔄 Open

Base: main ← Head: fix/flash-decode-singleton-strides

📝 Commits (1)

0632722 Fix singleton strides in CUDA flash decode

📊 Changes

2 files changed (+41 additions, -2 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+3 -2)
📝 tests/test_flash_attention.py (+38 -0)

📄 Description

PyTorch 2.8.0 exports singleton dimensions through DLPack with stride 1. The CUDA GQA decode and tree-merge checks reject these valid output layouts, even though singleton dimensions are only indexed at zero.

Check GQA output row strides only when query_length > 1, batch strides only when batch > 1, and tree-merge output row strides only when rows > 1.

Last-dimension contiguity, Q/K/V validation, and existing error messages are preserved. Kernel code and launch arguments are unchanged. The step-merge API benefits through its existing delegation to tree merge.

Validation on Windows with an RTX 4090:

With the full patch applied, tests/test_flash_attention.py:

PyTorch 2.8.0+cu128: 162 passed, 20 skipped.
PyTorch 2.14.1+cu130: 162 passed, 20 skipped.
Both runs include all 8 new regression cases. The 20 skipped tests are HIP/ROCm-specific.

Separately, only the 8 new regression cases were run on PyTorch 2.8.0+cu128, rebuilding the CUDA extension and starting a fresh Python process for each stage:

Original checks: 6 failed with Unsupported Flash Attention tensor strides, 2 passed.
GQA fix only: 4 failed with Unsupported tree merge strides, 4 passed.
Full patch restored: 8 passed.

None of these 8 cases were skipped. The intermediate stage independently exercises the tree-merge stride check.

Refs Comfy-Org/ComfyUI#17278.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/comfy-kitchen/pull/269 **Author:** [@curealb](https://github.com/curealb) **Created:** 10/22/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/flash-decode-singleton-strides` --- ### 📝 Commits (1) - [`0632722`](https://github.com/Comfy-Org/comfy-kitchen/commit/0632722d80227e876289f2bba61ca87d0f8ab01a) Fix singleton strides in CUDA flash decode ### 📊 Changes **2 files changed** (+41 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+3 -2) 📝 `tests/test_flash_attention.py` (+38 -0) </details> ### 📄 Description PyTorch 2.8.0 exports singleton dimensions through DLPack with stride 1. The CUDA GQA decode and tree-merge checks reject these valid output layouts, even though singleton dimensions are only indexed at zero. Check GQA output row strides only when query_length > 1, batch strides only when batch > 1, and tree-merge output row strides only when rows > 1. Last-dimension contiguity, Q/K/V validation, and existing error messages are preserved. Kernel code and launch arguments are unchanged. The step-merge API benefits through its existing delegation to tree merge. ## **Validation on Windows with an RTX 4090:** With the full patch applied, `tests/test_flash_attention.py`: - PyTorch 2.8.0+cu128: 162 passed, 20 skipped. - PyTorch 2.14.1+cu130: 162 passed, 20 skipped. Both runs include all 8 new regression cases. The 20 skipped tests are HIP/ROCm-specific. Separately, only the 8 new regression cases were run on PyTorch 2.8.0+cu128, rebuilding the CUDA extension and starting a fresh Python process for each stage: - Original checks: 6 failed with `Unsupported Flash Attention tensor strides`, 2 passed. - GQA fix only: 4 failed with `Unsupported tree merge strides`, 4 passed. - Full patch restored: 8 passed. None of these 8 cases were skipped. The intermediate stage independently exercises the tree-merge stride check. Refs Comfy-Org/ComfyUI#17278. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 22:22:20 +00:00

bsakaguchi referenced this issue

2026-10-26 22:22:23 +00:00

[PR #279] Accept DLPack-normalized size-1 strides in flash decode GQA #279