Comfy-Org / Comfy-Org/comfy-kitchen

Feature Request: Streamed QKV Projection and Packed Low-Bit K/V Support for Extreme Long-Sequence Attention

Open
#119 1 comment 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
220
Forks
91
Avg merge
1d 7h
Merged PRs (30d)
12

Description

Hi, I would like to propose a low-level attention/matmul feature that could be useful for MiniMax H3 and potentially other very long-sequence diffusion transformers.

Native Sol Attention now provides substantial compute speedups at high token counts, but there is still a separate **peak-memory problem before attention begins**: the QKV projection itself can be too large to materialize.

## Example: MiniMax H3

For an H3 Ref2VA workload with approximately **152,263 packed tokens**, I encountered:

```text
Currently allocated : 6.09 GiB
Requested : 6.11 GiB
Device limit : 11.94 GiB
Free : 0 bytes
```

The requested ~6.11 GiB is almost exactly the expected size of H3's full BF16 fused-QKV tensor at this sequence length.

Therefore, even a very memory-efficient attention kernel cannot help if:

```text
qkv = qkv_proj(x)
```

already needs to allocate ~6 GiB before the kernel receives Q/K/V.

## Proposed backend capability

It would be useful to support a **streamed QKV projection** that never materializes the complete high-precision fused-QKV output.

Conceptually:

```text
hidden states

streamed / tiled QKV GEMM

├── Q → BF16/FP16 storage

├── K → rotate + quantize → packed 4-bit storage

└── V → rotate + quantize → packed 4-bit storage
```

Attention would then directly consume:

```text
Q: BF16/FP16
K: packed INT4
V: packed INT4
```

and decode K/V only tile-by-tile inside the kernel.

This is different from adding a quantization operation after `qkv_proj`, because:

```text
full BF16 QKV
→ quantize
```

still requires the original peak allocation and therefore does not solve this class of OOM.

## Potential reduction

For the H3 example:

```text
BF16 QKV ≈ 6.1 GiB
```

Keeping Q in BF16 while storing K/V at 4-bit gives a theoretical payload around:

```text
Q BF16 ≈ 2.04 GiB
K 4-bit ≈ 0.51 GiB
V 4-bit ≈ 0.51 GiB

≈ 3.06 GiB + quantization metadata
```

So this could approximately halve the QKV working-set size.

## Quality-oriented low-bit format

I would not suggest naive INT4 as the final goal.

A more interesting starting point may be **fixed-width rotated INT4**, inspired by techniques such as TurboQuant:

* orthogonal / Hadamard-style rotation to suppress outliers
* per-group or per-tile scaling
* Q retained at BF16/FP16
* K/V compressed more aggressively
* optional bias/error correction if useful

TurboQuant demonstrated near-quality-neutral very-low-bit K/V compression in its evaluated LLM workloads.

HyperQuant also demonstrated that aggressive rotation-based low-bit quantization can remain highly accurate in large diffusion/video models such as LTX-2, although the LTX-2 experiment quantized linear weights rather than K/V activations.

Therefore, H3 K/V quality would still need direct validation rather than assuming the same near-lossless behavior.

For a first GPU implementation, fixed-width INT4 may be much more practical than variable-length Rice/entropy coding because Triton kernels can address packed tiles directly.

## Mixed precision for conditioning

For diffusion/video architectures with packed multimodal sequences, it may be useful to support different precision by token range.

For example, in H3:

```text
text conditioning → BF16 / INT8 K/V
reference images/videos → BF16 / INT8 K/V
audio conditioning → BF16 / INT8 K/V
generated video tokens → rotated INT4 K/V
```

This could protect prompt/reference fidelity while compressing the much larger generated-video portion.

Similar policies could also be applied by Transformer block or diffusion timestep.

## Dense fallback is important

Packed low-bit K/V support would ideally not exist only inside Sol Attention.

If a dense fallback performs:

```text
packed K/V
→ full BF16 expansion
→ dense attention
```

the original peak-memory problem can return.

A useful backend would therefore allow both sparse/Sol and dense attention paths to read packed K/V tile-by-tile without globally materializing BF16 K/V.

## Why this may be useful beyond H3

H3 is an extreme example because it combines:

* long target video sequences
* reference images
* reference videos
* audio
* text

into one packed sequence.

But the same issue can appear in future long-context video DiTs and multimodal diffusion transformers.

So a reusable capability such as:

```text
streamed_qkv_projection
+
packed low-bit K/V representation
+
packed-K/V dense/Sol attention
```

could potentially be useful as a general Comfy Kitchen primitive rather than an H3-only implementation.

The main goal would be **peak-VRAM reduction**, not merely attention speed.

Would this kind of streamed-QKV / packed-KV backend be within scope for Comfy Kitchen?

Contributor guide

Open the contributing guide

Research direction

No repository files or tests are identified in the request. Start by locating the existing Sol Attention backend and QKV projection entry point, then determine whether packed K/V can be consumed by both sparse/Sol and dense paths. Done would require an agreed scope plus direct H3 quality and peak-VRAM validation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.