Comfy-Org/comfy-kitchen

Fork 0

mirror of https://github.com/Comfy-Org/comfy-kitchen.git synced 2026-11-04 23:37:42 +00:00

[PR #236] Validate sol_attn binding arguments on CUDA like the HIP binding does #245

New issue

Open

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

bsakaguchi commented

2026-10-26 22:21:57 +00:00

Owner

📋 Pull Request Information

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

Base: main ← Head: sol-attn-cuda-binding-checks

📝 Commits (1)

3099629 Validate sol_attn binding arguments on CUDA like the HIP binding does

📊 Changes

1 file changed (+68 additions, -25 deletions)

View changed files

📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+68 -25)

📄 Description

Problem

tests/test_sol_attn.py::test_bindings_* (12 cases) fail on the CUDA backend with DID NOT RAISE RuntimeError. The CUDA _C.sol_attn / sol_attn_core / sol_producer_begin / sol_producer_chunk bindings check element counts and the q/k/v staging layout, but not:

the dtype or contiguity of threshold, key_bias, out, the producer operands and the stats buffers — a bfloat16 threshold of the right count, a strided threshold view, a float16 key_bias or an int8 out all get through,
that the workspace is uint8 (the plan offsets are bytes, and the Python layer slices the workspace with them),
non-positive batch / seq_len / num_heads (they size the workspace plan and every grid),
negative or inverted [start, end) sink ranges (a negative start makes the routing kernel emit extra entries as wrapped uint16 block ids and overrun the per-query row).

Because the kernels take bare pointers sized from (B, T, H), each of these is an out-of-bounds device access rather than an exception, followed by a sticky CUDA error for the rest of the process — in the test suite the first such case cascades into hundreds of torch.AcceleratorError failures. The HIP binding (backends/hip/dlpack_bindings.cpp) already rejects all of them with the messages the tests match on.

Change

One file, backends/cuda/dlpack_bindings.cpp, mirroring the HIP binding's helpers:

need_elems takes the expected map_dtype_to_code value and also requires a packed layout (reusing need_contiguous, moved above it).
need_extents and need_sinks added; need_workspace additionally requires uint8.
sol_attn, sol_attn_core, sol_producer_begin and sol_producer_chunk call them in the same order as the HIP binding. The sol_attn_core stats check becomes three need_elems calls so the message names the offending buffer.

Dtype codes per operand match what backends/cuda/__init__.py already passes (threshold/key_bias/fab/kmean/vscale/stats fp32, qkv/qw/kw/out bf16, workspace uint8), so production callers are unaffected.

Tests

A6000, sm86 build: tests/test_sol_attn.py 105 passed (before: 12 failed, plus the cascade when run without -x). tests/test_int8_attention.py unchanged (the two long_sequence_and_partial_tile[scale<=0] failures are the pre-existing torch-reference-NaN cases).

_{🔄 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/236 **Author:** [@Many0therFunctions](https://github.com/Many0therFunctions) **Created:** 10/10/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `sol-attn-cuda-binding-checks` --- ### 📝 Commits (1) - [`3099629`](https://github.com/Comfy-Org/comfy-kitchen/commit/3099629dab289d2b4de5d6984dc7d0fad17d8b90) Validate sol_attn binding arguments on CUDA like the HIP binding does ### 📊 Changes **1 file changed** (+68 additions, -25 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+68 -25) </details> ### 📄 Description ## Problem `tests/test_sol_attn.py::test_bindings_*` (12 cases) fail on the CUDA backend with `DID NOT RAISE RuntimeError`. The CUDA `_C.sol_attn` / `sol_attn_core` / `sol_producer_begin` / `sol_producer_chunk` bindings check element counts and the q/k/v staging layout, but not: - the dtype or contiguity of `threshold`, `key_bias`, `out`, the producer operands and the stats buffers — a bfloat16 `threshold` of the right count, a strided `threshold` view, a float16 `key_bias` or an int8 `out` all get through, - that the workspace is uint8 (the plan offsets are bytes, and the Python layer slices the workspace with them), - non-positive `batch` / `seq_len` / `num_heads` (they size the workspace plan and every grid), - negative or inverted `[start, end)` sink ranges (a negative start makes the routing kernel emit extra entries as wrapped uint16 block ids and overrun the per-query row). Because the kernels take bare pointers sized from `(B, T, H)`, each of these is an out-of-bounds device access rather than an exception, followed by a sticky CUDA error for the rest of the process — in the test suite the first such case cascades into hundreds of `torch.AcceleratorError` failures. The HIP binding (`backends/hip/dlpack_bindings.cpp`) already rejects all of them with the messages the tests match on. ## Change One file, `backends/cuda/dlpack_bindings.cpp`, mirroring the HIP binding's helpers: - `need_elems` takes the expected `map_dtype_to_code` value and also requires a packed layout (reusing `need_contiguous`, moved above it). - `need_extents` and `need_sinks` added; `need_workspace` additionally requires uint8. - `sol_attn`, `sol_attn_core`, `sol_producer_begin` and `sol_producer_chunk` call them in the same order as the HIP binding. The `sol_attn_core` stats check becomes three `need_elems` calls so the message names the offending buffer. Dtype codes per operand match what `backends/cuda/__init__.py` already passes (`threshold`/`key_bias`/`fab`/`kmean`/`vscale`/stats fp32, `qkv`/`qw`/`kw`/`out` bf16, workspace uint8), so production callers are unaffected. ## Tests A6000, sm86 build: `tests/test_sol_attn.py` 105 passed (before: 12 failed, plus the cascade when run without `-x`). `tests/test_int8_attention.py` unchanged (the two `long_sequence_and_partial_tile[scale<=0]` failures are the pre-existing torch-reference-NaN cases). --- <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:21:57 +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#245

No description provided.

Rows
Columns