Fix TritonVAE device mismatch in int8 conv path during lowvram/partial VAE loads - #688
Open
PlagueKind wants to merge 1 commit into
Open
Fix TritonVAE device mismatch in int8 conv path during lowvram/partial VAE loads#688PlagueKind wants to merge 1 commit into
PlagueKind wants to merge 1 commit into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Int8CausalConv3d,Int8Conv2d, andInt8InnerConv3dinnodes/triton_vae.pythrow a device-mismatch error (and, with--fast/certain triton versions, aValueError: Pointer argument cannot be accessed from Triton) when the VAE falls back to a partial/lowvram load instead of a full load.Repro: run VAEDecode (or an upscale-pass decode) after enough VRAM pressure that the VideoVAE gets staged as
loaded partially; 0.00 MB usable, ..., lowvram patches: 0instead of a full load. BothRuntimeError: Expected all tensors to be on the same deviceand, once the first issue is worked around,ValueError: Pointer argument cannot be accessed from Triton (cpu tensor?)are reproducible this way.Root cause
ComfyUI's dynamic VRAM / lowvram loader only device-casts real model parameters/buffers routed through
comfy.ops(i.e.weight). Two things in these int8 conv classes fall outside that path:int8_qw/int8_sw— registered viaregister_buffer(..., persistent=False), which is correct forstate_dictbehavior but doesn't put them in ComfyUI's lowvram casting sweep.self.bias— assigned directly from the original conv'sParameter(self.bias = orig.bias), sharing the same tensor object rather than being reinitialized through the patcher, so it isn't touched either.Under a full load both land on
x.deviceincidentally. Under a partial/lowvram load, both can remain oncpuwhile the activation tensorxis oncuda, causing the failures above at thescale_vec = self.int8_sw * amax.float()andbias.float()call sites.This mirrors an existing guard already present in
FusedRMSSiLU/FusedGNSiLUin the same file (if gamma.device != x.device: gamma = gamma.to(x.device)) — this PR just extends the same pattern to the int8 conv classes, which didn't have it.Fix
In
Int8CausalConv3d._int8_forward,Int8Conv2d.forward, andInt8InnerConv3d.forward: before use, checkint8_sw/int8_qwandbiasagainstx.deviceand move them if needed, same as the existing fused-norm classes already do forgamma/weight/bias.(with
int8_qw/biasused in place ofself.int8_qw/self.biasin the subsequent kernel call)Testing
Reproduced on LTX2.3 decode (
Int8CausalConv3dpath) with dynamic VRAM enabled, forcing a partial VAE load via VRAM pressure from a prior sampling stage. Confirmed the error no longer occurs after the fix, decode completes normally.Int8Conv2dandInt8InnerConv3dweren't independently reproduced but have the identical pattern, so the same fix is applied preemptively.