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

[GH-ISSUE #119] Feature Request: Streamed QKV Projection and Packed Low-Bit K/V Support for Extreme Long-Sequence Attention #37

New issue

Open

opened 2026-10-26 22:19:33 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

2026-10-26 22:19:33 +00:00

Owner

Originally created by @eliowang6 on GitHub (Aug 16, 2026).
Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/119

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:

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:

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:

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:

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:

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:

BF16 QKV ≈ 6.1 GiB

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

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 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:

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:

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?

Originally created by @eliowang6 on GitHub (Aug 16, 2026). Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/119 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?

bsakaguchi commented

2026-10-26 22:19:33 +00:00

Author

Owner

@coderabbitai[bot] commented on GitHub (Aug 16, 2026):

🔗 Related PRs

Comfy-Org/comfy-kitchen#94 - HIP: fix RDNA4 GEMM tile selection, vectorize fp8 elementwise kernels [open]
Comfy-Org/comfy-kitchen#103 - Implement optimized int8 attention (CORE-74) [merged]
Comfy-Org/comfy-kitchen#104 - HIP port of #103: optimized int8 attention [merged]
Comfy-Org/comfy-kitchen#108 - Optimize HIP ConvRot and deep-K INT8 kernels [open]

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord!

@coderabbitai[bot] commented on GitHub (Aug 16, 2026):  <details> <summary>🔗 Related PRs</summary> Comfy-Org/comfy-kitchen#94 - HIP: fix RDNA4 GEMM tile selection, vectorize fp8 elementwise kernels [open] Comfy-Org/comfy-kitchen#103 - Implement optimized int8 attention (CORE-74) [merged] Comfy-Org/comfy-kitchen#104 - HIP port of `#103`: optimized int8 attention [merged] Comfy-Org/comfy-kitchen#108 - Optimize HIP ConvRot and deep-K INT8 kernels [open] </details> --- <details> <summary> 🧪 Issue enrichment is currently in open beta.</summary> You can configure auto-planning by selecting labels in the issue_enrichment configuration. To disable automatic issue enrichment, add the following to your `.coderabbit.yaml`: ```yaml issue_enrichment: auto_enrich: enabled: false ``` </details> 💬 Have feedback or questions? Drop into our [discord](https://discord.gg/coderabbit)!

bsakaguchi referenced this issue

2026-10-26 22:20:19 +00:00

[PR #37] [MERGED] Add support for int8 quantization backend #94