You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Proposal: Generic memory-aware streaming execution for transformer models
Summary
I would like to propose moving a class of activation-memory optimizations currently implemented specifically for MiniMax H3 into generic ComfyUI infrastructure.
The basic idea is to add an inference-time streaming planner in Comfy core. The planner would inspect the final effective model and its patches immediately before inference, recognize transformer structures for which bounded execution is provably safe, query the selected compute/attention consumers for their streaming capabilities, and construct a lower-memory execution strategy.
The initial targets would be:
bounded/chunked MLP execution;
feature-sliced gated MLP execution where mathematically legal;
bounded Q/K/V projection;
retaining global K/V while streaming bounded Q;
streaming attention output directly through the output projection;
preserving Comfy's existing weight loading, patching, quantization and offloading behavior.
This would not be a sparse-attention proposal and would not intentionally change model semantics. It is about reducing peak activation memory by changing the execution schedule of mathematically equivalent operations.
Large video diffusion transformers can use several gigabytes of temporary memory in operations whose full intermediate tensors do not actually need to exist simultaneously.
In H3-Optimizations we currently reduce this by manually scheduling parts of MiniMax H3:
QKV projection is performed in bounded token slabs.
Compatible attention implementations retain global K/V while consuming bounded Q.
Attention output can be projected and written back in bounded slabs instead of retaining a full attention-output tensor.
MLPs operate on bounded token ranges.
ConvRot gated MLPs can additionally be divided along the FFN feature dimension so the full SwiGLU intermediate never exists.
Final projection/modulation is similarly bounded.
The implementation is currently MiniMax-specific because a custom node has to understand both the model and every attention consumer it wants to support.
The results suggest that the underlying idea is more general. In the H3 benchmark currently documented in the repository, adding chunked QKV, MLP and FinalLayer execution reduced peak VRAM by roughly 1.8 GB at 5 seconds and allowed the 10-second SageAttention workload to complete with roughly 2.7 GB less VRAM, with no measurable step-time regression attributable to those memory optimizations.
More importantly, most of the actual memory-saving mechanisms are not inherently specific to H3. What is H3-specific today is primarily the knowledge used to decide that those transformations are legal.
Proposed architecture
The proposal is to separate this into three responsibilities:
1. Comfy understands the effective model structure.
Comfy examines the final patched model and identifies execution patterns and dependencies.
2. Compute/attention consumers advertise what they can consume.
SDPA, Comfy Kitchen and third-party attention implementations explicitly declare their streaming capabilities rather than having the planner contain consumer-specific assumptions.
3. The planner finds the lowest-memory legal execution schedule.
It intersects model semantics, active patches, weight-format capabilities, consumer capabilities and hardware/backend constraints.
This turns the problem from “H3 knows how to invoke Kitchen/Sage/Triton” into “the model describes what may be decomposed and the consumer describes what it can accept.”
When planning should happen
Planning should occur against the final effective ModelPatcher state, after workflow nodes have had the opportunity to install model/object/attention patches and before denoiser execution begins.
The important point is not that every patched weight has already been physically materialized. Weight patches should remain under Comfy's normal lazy weight system.
The planner needs the final structural and execution state:
effective module structure;
object/forward replacements;
relevant transformer patches;
selected attention implementation;
tensor dimensions and head geometry;
quantized weight layouts;
hardware/backend capabilities.
It should plan how weights will be evaluated, not capture permanent copies of the weights themselves.
Actual weights should still be acquired through the normal Comfy mechanisms, including cast_bias_weight/uncast_bias_weight, so LoRAs, low-VRAM patches, DynamicVRAM/AIMDO and other weight transformations retain their existing behavior.
Model inspection
The planner should not try to reverse-engineer arbitrary Python.
Instead, it should recognize a conservative collection of common transformer structures and only optimize cases whose decomposition can be established safely.
For example, an MLP may be recognized as:
input projection;
known activation or gated activation;
output projection.
An attention block may expose or allow Comfy to infer:
separate or fused Q/K/V projections;
Q, K and V output ranges;
Q-head and KV-head counts;
head dimension;
optional Q/K normalization;
positional encoding/RoPE step;
attention consumer;
output projection.
This can be supplemented by lightweight model metadata where structural inspection alone cannot establish semantics.
There should be no need for checks equivalent to “if model is MiniMax H3.” H3 may provide enough metadata to describe its attention semantics, but the execution machinery should operate on generic descriptions.
MLP planning
MLPs are probably the easiest first target.
Basic token chunking is broadly applicable: instead of evaluating a huge MLP for every token simultaneously, the planner can process bounded token ranges while preserving the original residual/modulation semantics.
For gated MLPs, there is an additional generic optimization.
A gated MLP followed by a linear down projection can be partitioned along its FFN feature dimension. Each feature slice can be projected, activated, passed through the matching slice of the down projection, and accumulated into the output.
This allows execution without materializing the complete FFN-width activation.
Our H3 implementation currently uses this for ConvRot-256 SwiGLU MLPs, but the mathematical property is not H3-specific. The model-specific questions are simply:
what activation is being used;
how gate/up features correspond;
whether the output projection is linearly separable along that dimension;
what feature alignment the weight format requires.
Attention requires a stronger contract because different attention implementations have different input requirements.
The generic planner should be able to distinguish between at least:
consumers requiring complete Q/K/V;
consumers that can retain complete K/V while consuming Q incrementally;
consumers that can additionally use specialized K/V carriers;
consumers that can stream their output;
consumers with query/KV alignment requirements;
consumers with restrictions around masks, causal attention or GQA.
For a consumer supporting global K/V plus streamed Q, a legal plan can avoid retaining full Q:
project K/V using bounded source projection;
retain K/V for the attention invocation;
project one bounded Q region;
apply the model's Q normalization/positional transform;
execute attention for that Q region;
immediately run the corresponding output-projection region;
release the Q and attention-output intermediates;
repeat.
The source projection may remain BF16, FP16, FP8, ConvRot INT8, W4A8, or another format if that format provides the required projection/slicing operations. The planner should not require conversion to a preferred precision.
This is the part I think is most important to get right upstream.
Comfy should define an explicit streaming-consumer contract rather than maintaining a list of known implementations inside the planner.
A consumer should be able to advertise things such as:
whether bounded Q is accepted;
whether K must be global;
whether V must be global;
whether output can be produced in Q-sized slabs;
supported input/output dtypes;
Q/KV alignment requirements;
GQA support;
causal support;
mask/bias capabilities;
any specialized carrier requirements.
Comfy-owned consumers could then implement this contract directly.
PyTorch SDPA
SDPA would be a useful baseline implementation.
It can naturally support the common case of:
global K/V;
bounded Q;
bounded output.
The Comfy adapter would own the details of slicing masks correctly, accounting for query offsets where necessary, handling GQA and preserving the normal SDPA backend selection behavior.
This would provide generic memory-efficient execution without any custom attention kernel.
Comfy Kitchen
Comfy Kitchen could expose a stronger implementation where supported:
native quantized projection/carriers;
backend-specific Q/KV tile requirements;
global quantized K/V where applicable;
bounded Q;
bounded output.
Kitchen would own those implementation details. The generic Comfy planner would only consume its capability description.
Third-party consumers
SageAttention, custom sparse attention implementations, Sol Attention, etc. would not need Comfy to explicitly know about them.
If they want streamed execution, they implement the consumer contract.
If they do not implement it, Comfy treats them conservatively as requiring the traditional complete input tensors.
The proposed Comfy contract would generalize this and make the consumer responsible for supporting it.
Patch compatibility
Transformer patches should participate in the same capability system.
A patch that modifies Q/K/V or attention output may be compatible with streaming, or it may require complete tensors.
Ideally a patch could explicitly declare those constraints.
If compatibility is unknown, the planner should fail closed for the affected transformation and preserve current execution.
Conceptually, the legal plan is the intersection of:
model semantics ∩ active patches ∩ weight-format capabilities ∩ attention-consumer capabilities ∩ hardware constraints
No transformation should be selected unless all relevant components permit it.
This is preferable to trying to maintain recognition code for every custom node.
Weight-format capabilities
Quantized weight layouts could similarly expose whether particular decomposition operations are valid.
Useful generic capabilities include:
output-feature slicing;
input-feature slicing;
required slice alignment;
held-weight execution;
native activation fusion.
For example, ConvRot input-feature slicing needs to respect its rotation-group boundaries, while output-feature slicing needs to preserve the corresponding scale metadata.
These are properties of the weight layout, not H3.
Some of this may appropriately live in Comfy Kitchen, while the planner itself belongs in ComfyUI core.
Automatic chunk sizing
The first implementation could use conservative fixed chunk sizes, but the planner eventually has enough information to choose them automatically.
It knows:
sequence size;
hidden width;
FFN width;
activation dtype;
weight representation;
consumer tile/alignment requirements;
current device;
available working memory.
It could therefore estimate candidate execution schedules and select the largest efficient chunks that fit within a desired memory envelope.
This could eventually choose between strategies such as:
larger token chunks with feature-sliced MLP;
smaller token chunks with full-width MLP;
complete Q where cheap;
streamed Q where sequence length makes it worthwhile.
This is optional for an initial implementation; the key proposal is the contract/planner architecture.
Safety and fallback behavior
The intended default behavior should be conservative.
The planner should not silently quantize weights simply because a faster representation exists. Preserving the checkpoint's current representation should be the baseline behavior.
The transformations proposed here should preserve model computation within the normal numerical variation of changing kernel decomposition. Sparse attention, token dropping, approximation and other quality/performance tradeoffs are explicitly outside this proposal.
Suggested initial scope
I would suggest proving the architecture with a deliberately narrow first implementation:
Generic bounded MLP execution for recognized Linear → activation/gated activation → Linear structures.
Generic feature-sliced execution for one or two common gated-MLP forms.
Generic fused/separate QKV recognition for standard self-attention.
Global K/V + streamed Q/output as the first attention-streaming contract.
PyTorch SDPA as the baseline consumer.
Comfy Kitchen as the optimized/quantized consumer.
MiniMax H3 as the first large-model integration and regression target.
Flux would be a useful subsequent test because its DoubleStream blocks contain conventional fused QKV/MLPs while its SingleStream architecture combines QKV and MLP regions into shared projections. If the same planning model can eventually represent both H3 and Flux without model-name-specific execution code, that would be a strong indication that the abstraction is useful.
Why this belongs in Comfy core
A custom node can implement the memory optimizations, but it is in the wrong architectural position to make them generic.
Comfy core:
owns ModelPatcher and the final effective patch state;
owns weight acquisition/offloading;
owns built-in model implementations;
owns the built-in attention-selection machinery;
can provide contracts that third-party attention implementations can target;
can make SDPA and Comfy Kitchen reference implementations of those contracts.
A custom node instead has to learn model internals, identify every possible consumer itself, and continuously adapt to third-party ABI changes.
The H3 implementation has effectively reached the point where much of its complexity is compatibility glue rather than the underlying optimization.
The proposal is to move ownership of that compatibility to the components that actually know the relevant information.
The repository should be treated as evidence that the execution strategies are practical, not as a proposed API or code drop for Comfy. A core implementation should probably be significantly smaller and more generic because Comfy has information and lifecycle control that a custom node does not.
Questions for Comfy Org
Before attempting an implementation, I would mainly like feedback on the direction:
Is a generic inference-time activation/streaming planner something ComfyUI would want in core?
Does the proposed separation between model semantics, patch constraints, weight-format capabilities and attention-consumer capabilities fit Comfy's preferred architecture?
Would you prefer models to expose small explicit semantic descriptors, rely primarily on standardized module structures, or some combination?
Does an attention streaming contract belong alongside the current attention backend registration/selection infrastructure?
Would SDPA + Comfy Kitchen + MiniMax H3 be a reasonable first implementation scope?
Should this initially be opt-in/experimental, or is conservative automatic planning acceptable if every unknown case fails back to current execution?
If the general direction is wanted, where would you prefer the planner and contracts to live before an implementation PR is started?
The main thing I would like to establish first is whether Comfy Org sees value in owning this generically. If so, I can turn the existing H3 work into a more concrete design around whatever interfaces you would prefer.
reacted with thumbs up emoji reacted with thumbs down emoji reacted with laugh emoji reacted with hooray emoji reacted with confused emoji reacted with heart emoji reacted with rocket emoji reacted with eyes emoji
Uh oh!
There was an error while loading. Please reload this page.
Proposal: Generic memory-aware streaming execution for transformer models
Summary
I would like to propose moving a class of activation-memory optimizations currently implemented specifically for MiniMax H3 into generic ComfyUI infrastructure.
The basic idea is to add an inference-time streaming planner in Comfy core. The planner would inspect the final effective model and its patches immediately before inference, recognize transformer structures for which bounded execution is provably safe, query the selected compute/attention consumers for their streaming capabilities, and construct a lower-memory execution strategy.
The initial targets would be:
This would not be a sparse-attention proposal and would not intentionally change model semantics. It is about reducing peak activation memory by changing the execution schedule of mathematically equivalent operations.
The existing [H3-Optimizations](https://github.com/Zironic/H3-Optimizations) custom node is intended as prior art and a working prototype of many of these techniques.
Motivation
Large video diffusion transformers can use several gigabytes of temporary memory in operations whose full intermediate tensors do not actually need to exist simultaneously.
In H3-Optimizations we currently reduce this by manually scheduling parts of MiniMax H3:
The implementation is currently MiniMax-specific because a custom node has to understand both the model and every attention consumer it wants to support.
The results suggest that the underlying idea is more general. In the H3 benchmark currently documented in the repository, adding chunked QKV, MLP and FinalLayer execution reduced peak VRAM by roughly 1.8 GB at 5 seconds and allowed the 10-second SageAttention workload to complete with roughly 2.7 GB less VRAM, with no measurable step-time regression attributable to those memory optimizations.
More importantly, most of the actual memory-saving mechanisms are not inherently specific to H3. What is H3-specific today is primarily the knowledge used to decide that those transformations are legal.
Proposed architecture
The proposal is to separate this into three responsibilities:
1. Comfy understands the effective model structure.
Comfy examines the final patched model and identifies execution patterns and dependencies.
2. Compute/attention consumers advertise what they can consume.
SDPA, Comfy Kitchen and third-party attention implementations explicitly declare their streaming capabilities rather than having the planner contain consumer-specific assumptions.
3. The planner finds the lowest-memory legal execution schedule.
It intersects model semantics, active patches, weight-format capabilities, consumer capabilities and hardware/backend constraints.
This turns the problem from “H3 knows how to invoke Kitchen/Sage/Triton” into “the model describes what may be decomposed and the consumer describes what it can accept.”
When planning should happen
Planning should occur against the final effective ModelPatcher state, after workflow nodes have had the opportunity to install model/object/attention patches and before denoiser execution begins.
The important point is not that every patched weight has already been physically materialized. Weight patches should remain under Comfy's normal lazy weight system.
The planner needs the final structural and execution state:
It should plan how weights will be evaluated, not capture permanent copies of the weights themselves.
Actual weights should still be acquired through the normal Comfy mechanisms, including
cast_bias_weight/uncast_bias_weight, so LoRAs, low-VRAM patches, DynamicVRAM/AIMDO and other weight transformations retain their existing behavior.Model inspection
The planner should not try to reverse-engineer arbitrary Python.
Instead, it should recognize a conservative collection of common transformer structures and only optimize cases whose decomposition can be established safely.
For example, an MLP may be recognized as:
An attention block may expose or allow Comfy to infer:
This can be supplemented by lightweight model metadata where structural inspection alone cannot establish semantics.
There should be no need for checks equivalent to “if model is MiniMax H3.” H3 may provide enough metadata to describe its attention semantics, but the execution machinery should operate on generic descriptions.
MLP planning
MLPs are probably the easiest first target.
Basic token chunking is broadly applicable: instead of evaluating a huge MLP for every token simultaneously, the planner can process bounded token ranges while preserving the original residual/modulation semantics.
For gated MLPs, there is an additional generic optimization.
A gated MLP followed by a linear down projection can be partitioned along its FFN feature dimension. Each feature slice can be projected, activated, passed through the matching slice of the down projection, and accumulated into the output.
This allows execution without materializing the complete FFN-width activation.
Our H3 implementation currently uses this for ConvRot-256 SwiGLU MLPs, but the mathematical property is not H3-specific. The model-specific questions are simply:
The current implementation can be seen in [H3-Optimizations' bounded MLP machinery](https://github.com/Zironic/H3-Optimizations/tree/main/h3_optimizations/memory). The existing code also demonstrates safe fallback from specialized ConvRot execution to ordinary held/module execution when the specialized assumptions cannot be satisfied.
QKV and attention planning
Attention requires a stronger contract because different attention implementations have different input requirements.
The generic planner should be able to distinguish between at least:
For a consumer supporting global K/V plus streamed Q, a legal plan can avoid retaining full Q:
The source projection may remain BF16, FP16, FP8, ConvRot INT8, W4A8, or another format if that format provides the required projection/slicing operations. The planner should not require conversion to a preferred precision.
H3-Optimizations currently implements variants of this for several formats in its [QKV implementation](https://github.com/Zironic/H3-Optimizations/tree/main/h3_optimizations/qkv).
Attention consumer contract
This is the part I think is most important to get right upstream.
Comfy should define an explicit streaming-consumer contract rather than maintaining a list of known implementations inside the planner.
A consumer should be able to advertise things such as:
Comfy-owned consumers could then implement this contract directly.
PyTorch SDPA
SDPA would be a useful baseline implementation.
It can naturally support the common case of:
The Comfy adapter would own the details of slicing masks correctly, accounting for query offsets where necessary, handling GQA and preserving the normal SDPA backend selection behavior.
This would provide generic memory-efficient execution without any custom attention kernel.
Comfy Kitchen
Comfy Kitchen could expose a stronger implementation where supported:
Kitchen would own those implementation details. The generic Comfy planner would only consume its capability description.
Third-party consumers
SageAttention, custom sparse attention implementations, Sol Attention, etc. would not need Comfy to explicitly know about them.
If they want streamed execution, they implement the consumer contract.
If they do not implement it, Comfy treats them conservatively as requiring the traditional complete input tensors.
H3-Optimizations already contains a small proof of this model in its current [external streamed-attention consumer contract](https://github.com/Zironic/H3-Optimizations/blob/main/h3_optimizations/external_consumer.py). An override explicitly opts into streamed Q and receives a Q chunk against global K/V. Opaque overrides are left on the existing full-Q path rather than being guessed about.
The proposed Comfy contract would generalize this and make the consumer responsible for supporting it.
Patch compatibility
Transformer patches should participate in the same capability system.
A patch that modifies Q/K/V or attention output may be compatible with streaming, or it may require complete tensors.
Ideally a patch could explicitly declare those constraints.
If compatibility is unknown, the planner should fail closed for the affected transformation and preserve current execution.
Conceptually, the legal plan is the intersection of:
model semantics ∩ active patches ∩ weight-format capabilities ∩ attention-consumer capabilities ∩ hardware constraints
No transformation should be selected unless all relevant components permit it.
This is preferable to trying to maintain recognition code for every custom node.
Weight-format capabilities
Quantized weight layouts could similarly expose whether particular decomposition operations are valid.
Useful generic capabilities include:
For example, ConvRot input-feature slicing needs to respect its rotation-group boundaries, while output-feature slicing needs to preserve the corresponding scale metadata.
These are properties of the weight layout, not H3.
Some of this may appropriately live in Comfy Kitchen, while the planner itself belongs in ComfyUI core.
Automatic chunk sizing
The first implementation could use conservative fixed chunk sizes, but the planner eventually has enough information to choose them automatically.
It knows:
It could therefore estimate candidate execution schedules and select the largest efficient chunks that fit within a desired memory envelope.
This could eventually choose between strategies such as:
This is optional for an initial implementation; the key proposal is the contract/planner architecture.
Safety and fallback behavior
The intended default behavior should be conservative.
The planner should not silently quantize weights simply because a faster representation exists. Preserving the checkpoint's current representation should be the baseline behavior.
The transformations proposed here should preserve model computation within the normal numerical variation of changing kernel decomposition. Sparse attention, token dropping, approximation and other quality/performance tradeoffs are explicitly outside this proposal.
Suggested initial scope
I would suggest proving the architecture with a deliberately narrow first implementation:
Flux would be a useful subsequent test because its DoubleStream blocks contain conventional fused QKV/MLPs while its SingleStream architecture combines QKV and MLP regions into shared projections. If the same planning model can eventually represent both H3 and Flux without model-name-specific execution code, that would be a strong indication that the abstraction is useful.
Why this belongs in Comfy core
A custom node can implement the memory optimizations, but it is in the wrong architectural position to make them generic.
Comfy core:
A custom node instead has to learn model internals, identify every possible consumer itself, and continuously adapt to third-party ABI changes.
The H3 implementation has effectively reached the point where much of its complexity is compatibility glue rather than the underlying optimization.
The proposal is to move ownership of that compatibility to the components that actually know the relevant information.
Prior art / prototype
The current prototype is:
[Zironic/H3-Optimizations](https://github.com/Zironic/H3-Optimizations)
Relevant areas include:
The repository should be treated as evidence that the execution strategies are practical, not as a proposed API or code drop for Comfy. A core implementation should probably be significantly smaller and more generic because Comfy has information and lifecycle control that a custom node does not.
Questions for Comfy Org
Before attempting an implementation, I would mainly like feedback on the direction:
The main thing I would like to establish first is whether Comfy Org sees value in owning this generically. If so, I can turn the existing H3 work into a more concrete design around whatever interfaces you would prefer.
All reactions