[PR #126] Add minimal TensorArg abstraction for strided RoPE kernels #161

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/126
Author: @contentis
Created: 8/21/2026
Status: 🔄 Open

Base: main ← Head: tensor-abstraction

📝 Commits (8)

0b58908 Add minimal TensorArg abstraction for strided RoPE kernels
099176a Fix and consolidate TensorArg bindings
305cc26 Fix nanobind compatibility and FP8 dtype mapping
29cfe23 Extend TensorArg across CUDA operations
50d6e6c Expand TensorArg CUDA coverage
af829de Complete CUDA TensorArg launcher migration
d2658bd Fix cuBLAS workspace and AWQ validation
1f0a06b Limit TensorArg adoption to metadata-aware kernels

📊 Changes

12 files changed (+508 additions, -520 deletions)

View changed files

📝 .github/workflows/build-wheels.yml (+4 -4)
📝 MANIFEST.in (+2 -0)
📝 comfy_kitchen/backends/cuda/CMakeLists.txt (+2 -0)
📝 comfy_kitchen/backends/cuda/dlpack_bindings.cpp (+137 -322)
📝 comfy_kitchen/backends/cuda/ops/flash_decode.cu (+34 -21)
📝 comfy_kitchen/backends/cuda/ops/rms_rope.cu (+125 -142)
📝 comfy_kitchen/backends/cuda/sage_attention/quant_qk_int8.cu (+20 -6)
📝 comfy_kitchen/backends/cuda/sage_attention/quant_v_int8.cu (+17 -10)
📝 comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu (+70 -13)
➕ comfy_kitchen/backends/tensor.h (+43 -0)
➕ comfy_kitchen/backends/tensor_bindings.h (+52 -0)
📝 pyproject.toml (+2 -2)

📄 Description

Introduces pointer-free TensorMeta and pass-by-value TensorArg descriptors. Migrates CUDA and HIP apply-RoPE/RMS-RoPE bindings and launchers away from expanded pointer, shape, stride, and dtype argument lists while preserving strided-view support.

This aims to improve readabilty - if this deisgn is acceptable we can apply it to the other kernels aswell.

_{🔄 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/126 **Author:** [@contentis](https://github.com/contentis) **Created:** 8/21/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `tensor-abstraction` --- ### 📝 Commits (8) - [`0b58908`](https://github.com/Comfy-Org/comfy-kitchen/commit/0b58908c0a8c0533203cc3da4dc263f3148805d8) Add minimal TensorArg abstraction for strided RoPE kernels - [`099176a`](https://github.com/Comfy-Org/comfy-kitchen/commit/099176a8ac7d7c70213d400ed3a58310e98c776c) Fix and consolidate TensorArg bindings - [`305cc26`](https://github.com/Comfy-Org/comfy-kitchen/commit/305cc268f286eb89840f50b27ef1f911d245f8ef) Fix nanobind compatibility and FP8 dtype mapping - [`29cfe23`](https://github.com/Comfy-Org/comfy-kitchen/commit/29cfe23323ef7facb54182c57548d473ce63613b) Extend TensorArg across CUDA operations - [`50d6e6c`](https://github.com/Comfy-Org/comfy-kitchen/commit/50d6e6c1e9191efceab6b5a8f334e7e4d92bb963) Expand TensorArg CUDA coverage - [`af829de`](https://github.com/Comfy-Org/comfy-kitchen/commit/af829de1bcda631e5194ae9b496cc059109fdc54) Complete CUDA TensorArg launcher migration - [`d2658bd`](https://github.com/Comfy-Org/comfy-kitchen/commit/d2658bdcb8f91261e3ec93b9aa5f7ff4a3a178f3) Fix cuBLAS workspace and AWQ validation - [`1f0a06b`](https://github.com/Comfy-Org/comfy-kitchen/commit/1f0a06b99e2befdf0533573e6bbcf6521440404d) Limit TensorArg adoption to metadata-aware kernels ### 📊 Changes **12 files changed** (+508 additions, -520 deletions) <details> <summary>View changed files</summary> 📝 `.github/workflows/build-wheels.yml` (+4 -4) 📝 `MANIFEST.in` (+2 -0) 📝 `comfy_kitchen/backends/cuda/CMakeLists.txt` (+2 -0) 📝 `comfy_kitchen/backends/cuda/dlpack_bindings.cpp` (+137 -322) 📝 `comfy_kitchen/backends/cuda/ops/flash_decode.cu` (+34 -21) 📝 `comfy_kitchen/backends/cuda/ops/rms_rope.cu` (+125 -142) 📝 `comfy_kitchen/backends/cuda/sage_attention/quant_qk_int8.cu` (+20 -6) 📝 `comfy_kitchen/backends/cuda/sage_attention/quant_v_int8.cu` (+17 -10) 📝 `comfy_kitchen/backends/cuda/sage_attention/sage_attn_launcher.cu` (+70 -13) ➕ `comfy_kitchen/backends/tensor.h` (+43 -0) ➕ `comfy_kitchen/backends/tensor_bindings.h` (+52 -0) 📝 `pyproject.toml` (+2 -2) </details> ### 📄 Description Introduces pointer-free TensorMeta<Rank> and pass-by-value TensorArg<Rank> descriptors. Migrates CUDA and HIP apply-RoPE/RMS-RoPE bindings and launchers away from expanded pointer, shape, stride, and dtype argument lists while preserving strided-view support. This aims to improve readabilty - if this deisgn is acceptable we can apply it to the other kernels aswell. --- <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:01 +00:00

bsakaguchi referenced this issue

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

[PR #161] [CLOSED] [Feat][Ascend] Add ConvRot W4A4 linear backend #186

bsakaguchi referenced this issue

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

[PR #199] [MERGED] [BugFix][Ascend] Gate RoPE tests on operator capabilities and clarify support #214

bsakaguchi referenced this issue

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