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 #180] Bound VSA coarse score memory on HIP #199

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/comfy-kitchen/pull/180
Author: @Siilent-Sea
Created: 9/16/2026
Status: 🔄 Open

Base: main ← Head: codex/vsa-coarse-query-chunk

📝 Commits (1)

88569c7 Bound VSA coarse score memory on HIP

📊 Changes

3 files changed (+48 additions, -6 deletions)

View changed files

📝 comfy_kitchen/backends/eager/sol_attn.py (+14 -4)
📝 comfy_kitchen/backends/hip/__init__.py (+4 -2)
➕ tests/test_eager_sol_attn.py (+30 -0)

📄 Description

Problem and scope

The VSA coarse branch currently materializes the complete [BH, N, N] FP32
score tensor. At the production-shape synthetic geometry used here
(BH=56, N=1338, D=128), that score path accounts for a large avoidable
peak allocation.

Why the change is general

Each query row is independent once the complete key/value block means are
available. Slicing only the query dimension preserves the full key set and the
per-query softmax denominator; it does not change VSA routing, keys, values, or
the mathematical operation.

Implementation

Add an optional query_chunk argument to eager coarse_output.
Preserve the existing dense implementation when it is omitted.
Preallocate the result and process query slices without retaining every
slice before concatenation.
Pass query_chunk=64 from the two HIP coarse-gate call sites only.

Regression test

The new CPU test compares dense and query-chunked results for both a divisible
and a ragged final chunk (131/64 and 17/7).

Numerical validation

Current main: 6c7f08bf8d8c6d039892b311d40c976c26ce979f
Deterministic seed: 271
Tolerance fixed before the target-device run: rtol=2e-4, atol=5e-5
Mac CPU smoke max absolute difference: 4.172325134277344e-07
gfx1201 max absolute difference: 5.066394805908203e-07

Performance and memory validation

Environment: WSL2 Ubuntu 24.04, AMD Radeon RX 9070 XT (gfx1201), PyTorch
2.9.1+rocm7.2.1.gitff65f5bc, HIP 7.2.53211-e1a6bc5663, FP32 synthetic
block-mean tensors of shape [56, 1338, 128], two warmups and seven measured
runs.

implementation	median peak allocation delta	median wall
current dense	801.4609 MiB	18.040 ms
query chunk 64	91.4648 MiB	17.519 ms

That is an 88.59% reduction in the measured peak allocation delta with no
timing regression in this isolated coarse-branch benchmark.

Raw warm wall samples (seconds)

dense: 0.018001765, 0.017981337, 0.018061866, 0.018013761, 0.018130983, 0.018103780, 0.018040032
chunked: 0.017591071, 0.017467413, 0.017297295, 0.017588049, 0.017518593, 0.017421424, 0.018291322

Platforms tested / not tested

Tested on Mac CPU for current-main import/numerics and on gfx1201 HIP for
numerics, peak allocation, and timing. This is not an H3 end-to-end benchmark.
CUDA/NVIDIA and other AMD generations were not measured.

Risks and fallback

The shared eager helper remains byte-for-byte on its existing dense branch when
query_chunk is omitted. The behavior change is limited to the two HIP
coarse-gate call sites; removing their argument restores the previous path.

Privacy and asset provenance

All validation uses deterministic synthetic tensors. No video, prompt, latent,
model/adapter weight, private path, or private hash is included.

Searches for VSA coarse attention, sol_attn memory, coarse_output, and
query chunk attention found no existing issue or PR implementing this bound.

Contribution status

The account holder will complete any platform-requested DCO or Comfy CLA step.

AI assistance disclosure

Submitted by the project owner with assistance from an owner-authorized AI
engineering agent. Reproduction, code, tests and benchmark artifacts are
included.

_{🔄 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/180 **Author:** [@Siilent-Sea](https://github.com/Siilent-Sea) **Created:** 9/16/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `codex/vsa-coarse-query-chunk` --- ### 📝 Commits (1) - [`88569c7`](https://github.com/Comfy-Org/comfy-kitchen/commit/88569c762f27d8de6473f9dc0d372ebe84515fa6) Bound VSA coarse score memory on HIP ### 📊 Changes **3 files changed** (+48 additions, -6 deletions) <details> <summary>View changed files</summary> 📝 `comfy_kitchen/backends/eager/sol_attn.py` (+14 -4) 📝 `comfy_kitchen/backends/hip/__init__.py` (+4 -2) ➕ `tests/test_eager_sol_attn.py` (+30 -0) </details> ### 📄 Description ## Problem and scope The VSA coarse branch currently materializes the complete `[BH, N, N]` FP32 score tensor. At the production-shape synthetic geometry used here (`BH=56`, `N=1338`, `D=128`), that score path accounts for a large avoidable peak allocation. ## Why the change is general Each query row is independent once the complete key/value block means are available. Slicing only the query dimension preserves the full key set and the per-query softmax denominator; it does not change VSA routing, keys, values, or the mathematical operation. ## Implementation - Add an optional `query_chunk` argument to eager `coarse_output`. - Preserve the existing dense implementation when it is omitted. - Preallocate the result and process query slices without retaining every slice before concatenation. - Pass `query_chunk=64` from the two HIP coarse-gate call sites only. ## Regression test The new CPU test compares dense and query-chunked results for both a divisible and a ragged final chunk (`131/64` and `17/7`). ## Numerical validation - Current main: `6c7f08bf8d8c6d039892b311d40c976c26ce979f` - Deterministic seed: `271` - Tolerance fixed before the target-device run: `rtol=2e-4`, `atol=5e-5` - Mac CPU smoke max absolute difference: `4.172325134277344e-07` - gfx1201 max absolute difference: `5.066394805908203e-07` ## Performance and memory validation Environment: WSL2 Ubuntu 24.04, AMD Radeon RX 9070 XT (gfx1201), PyTorch `2.9.1+rocm7.2.1.gitff65f5bc`, HIP `7.2.53211-e1a6bc5663`, FP32 synthetic block-mean tensors of shape `[56, 1338, 128]`, two warmups and seven measured runs. | implementation | median peak allocation delta | median wall | | --- | ---: | ---: | | current dense | 801.4609 MiB | 18.040 ms | | query chunk 64 | 91.4648 MiB | 17.519 ms | That is an 88.59% reduction in the measured peak allocation delta with no timing regression in this isolated coarse-branch benchmark. <details><summary>Raw warm wall samples (seconds)</summary> - dense: `0.018001765, 0.017981337, 0.018061866, 0.018013761, 0.018130983, 0.018103780, 0.018040032` - chunked: `0.017591071, 0.017467413, 0.017297295, 0.017588049, 0.017518593, 0.017421424, 0.018291322` </details> ## Platforms tested / not tested Tested on Mac CPU for current-main import/numerics and on gfx1201 HIP for numerics, peak allocation, and timing. This is not an H3 end-to-end benchmark. CUDA/NVIDIA and other AMD generations were not measured. ## Risks and fallback The shared eager helper remains byte-for-byte on its existing dense branch when `query_chunk` is omitted. The behavior change is limited to the two HIP coarse-gate call sites; removing their argument restores the previous path. ## Privacy and asset provenance All validation uses deterministic synthetic tensors. No video, prompt, latent, model/adapter weight, private path, or private hash is included. ## Related issue or discussion Searches for `VSA coarse attention`, `sol_attn memory`, `coarse_output`, and `query chunk attention` found no existing issue or PR implementing this bound. ## Contribution status The account holder will complete any platform-requested DCO or Comfy CLA step. ## AI assistance disclosure Submitted by the project owner with assistance from an owner-authorized AI engineering agent. Reproduction, code, tests and benchmark artifacts are included. --- <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:30 +00:00

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:41 +00:00

[PR #201] [MERGED] [Perf][Ascend] Avoid activation pack/unpack in ConvRot W4A4 #218