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

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/279
Author: @potoior
Created: 10/26/2026
Status: 🔄 Open

Base: main ← Head: dlpack-size1-stride-flash-decode

📝 Commits (1)

c41c5d5 Accept DLPack-normalized size-1 strides in flash decode GQA

📊 Changes

3 files changed (+47 additions, -2 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+8 -1)
📝 comfy_kitchen/backends/hip/dlpack_bindings.cpp (+7 -1)
📝 tests/test_flash_attention.py (+32 -0)

📄 Description

Fixes #274

Problem

flash_attention_decode_gqa rejects the first decode step of a single-token generate: PyTorch 2.9-era DLPack exports (observed on the NGC 25.09 container, GB10) normalize the strides of size-1 dimensions to 1, so the output [B, 1, H*D] reaches the binding with strides (1, 1, 1) and the binding's exact stride assertions throw Unsupported Flash Attention tensor strides (#274). The wrapper allocates that output itself, so nothing on the Python side can work around it.

Change

Size-1 strides never participate in addressing (the index is always 0), so the binding now only constrains the strides of non-singleton dimensions, in both the CUDA and HIP bindings. The kernel keeps receiving whatever stride the export produced: for every dimension the assertions previously guarded, either the dimension is larger than 1 and the stride is still checked, or the dimension is 1 and the stride value is irrelevant to addressing.

Tests

test_flash_attention_decode_gqa_accepts_dlpack_normalized_size1_strides drives the CUDA binding directly with an as_strided output whose strides mirror the normalized export — rejected before this change, accepted and numerically correct (against the softmax reference) after it. The test does not depend on the exporter's behavior, so it reproduces the #274 conditions on any torch.

Verified on an NVIDIA Thor (sm_110, torch 2.11.0+cu130, whose DLPack export still preserves size-1 strides): the new test plus the full test_flash_attention.py suite pass (155 passed, 20 HIP-only skips). The HIP binding change mirrors the CUDA one line for line but could not be exercised without AMD hardware.

_{🔄 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/279 **Author:** [@potoior](https://github.com/potoior) **Created:** 10/26/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `dlpack-size1-stride-flash-decode` --- ### 📝 Commits (1) - [`c41c5d5`](https://github.com/Comfy-Org/comfy-kitchen/commit/c41c5d5aa228d14ac120164f59b8cc33b274108a) Accept DLPack-normalized size-1 strides in flash decode GQA ### 📊 Changes **3 files changed** (+47 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+8 -1) 📝 `comfy_kitchen/backends/hip/dlpack_bindings.cpp` (+7 -1) 📝 `tests/test_flash_attention.py` (+32 -0) </details> ### 📄 Description Fixes #274 **Problem** `flash_attention_decode_gqa` rejects the first decode step of a single-token generate: PyTorch 2.9-era DLPack exports (observed on the NGC 25.09 container, GB10) normalize the strides of size-1 dimensions to 1, so the output `[B, 1, H*D]` reaches the binding with strides `(1, 1, 1)` and the binding's exact stride assertions throw `Unsupported Flash Attention tensor strides` (#274). The wrapper allocates that output itself, so nothing on the Python side can work around it. **Change** Size-1 strides never participate in addressing (the index is always 0), so the binding now only constrains the strides of non-singleton dimensions, in both the CUDA and HIP bindings. The kernel keeps receiving whatever stride the export produced: for every dimension the assertions previously guarded, either the dimension is larger than 1 and the stride is still checked, or the dimension is 1 and the stride value is irrelevant to addressing. **Tests** `test_flash_attention_decode_gqa_accepts_dlpack_normalized_size1_strides` drives the CUDA binding directly with an `as_strided` output whose strides mirror the normalized export — rejected before this change, accepted and numerically correct (against the softmax reference) after it. The test does not depend on the exporter's behavior, so it reproduces the #274 conditions on any torch. Verified on an NVIDIA Thor (sm_110, torch 2.11.0+cu130, whose DLPack export still preserves size-1 strides): the new test plus the full `test_flash_attention.py` suite pass (155 passed, 20 HIP-only skips). The HIP binding change mirrors the CUDA one line for line but could not be exercised without AMD hardware. --- <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:23 +00:00

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

Comfy-Org/comfy-kitchen#279

No description provided.

Rows
Columns