Skip to content

Fix TritonVAE device mismatch in int8 conv path during lowvram/partial VAE loads - #688

Open
PlagueKind wants to merge 1 commit into
kijai:mainfrom
PlagueKind:fix/tritonvae
Open

Fix TritonVAE device mismatch in int8 conv path during lowvram/partial VAE loads#688
PlagueKind wants to merge 1 commit into
kijai:mainfrom
PlagueKind:fix/tritonvae

Conversation

@PlagueKind

Copy link
Copy Markdown

Problem

Int8CausalConv3d, Int8Conv2d, and Int8InnerConv3d in nodes/triton_vae.py throw a device-mismatch error (and, with --fast/certain triton versions, a ValueError: 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: 0 instead of a full load. Both RuntimeError: Expected all tensors to be on the same device and, 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:

  1. int8_qw / int8_sw — registered via register_buffer(..., persistent=False), which is correct for state_dict behavior but doesn't put them in ComfyUI's lowvram casting sweep.
  2. self.bias — assigned directly from the original conv's Parameter (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.device incidentally. Under a partial/lowvram load, both can remain on cpu while the activation tensor x is on cuda, causing the failures above at the scale_vec = self.int8_sw * amax.float() and bias.float() call sites.

This mirrors an existing guard already present in FusedRMSSiLU/FusedGNSiLU in 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, and Int8InnerConv3d.forward: before use, check int8_sw/int8_qw and bias against x.device and move them if needed, same as the existing fused-norm classes already do for gamma/weight/bias.

int8_qw = self.int8_qw
int8_sw = self.int8_sw
if int8_sw.device != x.device:  # lowvram partial load keeps custom buffers off-device
    int8_qw = int8_qw.to(x.device)
    int8_sw = int8_sw.to(x.device)
scale_vec = int8_sw * amax.float()
bias = self.bias
if bias is not None and bias.device != x.device:  # lowvram partial load keeps bias param off-device
    bias = bias.to(x.device)

(with int8_qw/bias used in place of self.int8_qw/self.bias in the subsequent kernel call)

Testing

Reproduced on LTX2.3 decode (Int8CausalConv3d path) 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. Int8Conv2d and Int8InnerConv3d weren't independently reproduced but have the identical pattern, so the same fix is applied preemptively.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant