Blockswap (maybe) filling up the memory and the problem of Sageattn_3 working together with torch compile on wanvideowrapper1.4.7

Start by reproducing the reported memory growth with WanVideo Blockswap and the SageAttention 3 plus Torch Compile combination; the report gives the environment and workflow details but no standalone test or reproduction steps. Read custom_nodes\ComfyUI-WanVideoWrapper\wanvideo\modules\model.py and custom_nodes\ComfyUI-WanVideoWrapper\wanvideo\modules\attention.py, then compare the reported behavior with SageAttention 2.2. Done means identifying and documenting reproducible causes for the memory growth and compile failure without relying on the reporter's experimental edits.

My English is very poor. I wrote this article using a translator.

I'm not a developer. This is my first time posting issues. Please forgive me if there's anything inappropriate.I want to share my experiences with everyone.I'm not sure if it's appropriate to write here.

Platform：
windows 23H2，comfyui 0.17.2，WanVideoWrapper 1.4.7，python 3.12.0，pytorch 2.9.1，triton 3.5.1.post24、torchvision 0.24.1 、torchaudio 2.9.1，CUDA 13.0，NVIDIA Driver.580.97，5060ti 16G VRAM + 96G DRAM,
When installing these basic environments, I always specified the versions manually and did not let it choose automatically. Also, my comfyui is a portable version.

I have two main issues!

1. The WanVideo blockswap feature causes my memory usage to increase continuously. After generating videos multiple times, my 96GB memory would be fully occupied by the third round at most. To ensure continuous video generation, I used the --disable-pinned-memory command and an ancient standalone node called WanVideo Enhanced BlockSwap.
2. I attempted to enable the WanVideo Torch Compile Settings node in the Sage Attention 3.0 mode under the wanvideowrapper workflow by using the reduce-overhead parameter, in an effort to maximize the generation speed. However, this attempt failed.

I have to mention the long, messy, fumbling process I went through.

Regarding Question One：

---

No matter how much pinned memory I set in ComfyUI\comfy\model_management.py, it will eventually cause my 96G DRAM to be fully occupied. Only by turning off pinned memory can it function normally.
During the process of filling up the memory, it gradually increases with each round of video generation. There are no errors. It will keep filling up the physical memory until it is completely full. Then, Windows11 will automatically start using the virtual memory, and the speed will become extremely slow. I can only restart the comfyui. Clicking the "Clear Cache" button in the manager does not work.In my workflow, there are the LayerUtility: PurgeVRAM V2 node and the easy clearCacheAll node. However, neither of them works, even for the parts that exceed the pinned memory.
And the problem is not just like this. In order to improve the image quality as much as possible, I forcibly changed the attention mechanism in the first step to "SDPA" in ComfyUI-WanVideoWrapper\wanvideo\modules\model.py.
Perhaps due to the influence of the asynchronous transmission in the new version of ComfyUI, the computing efficiency of the graphics card is very low, and the load always fluctuates up and down on a 50% line.

If the resolution is increased, this sampling low-load process will last for an extremely long time. Not only for SDPA, but even when using the comfy attention mode, the same situation will occur.
As a comparison, I also kept an extremely stable old version of ComfyUI 0.03.39, in which I did not have this problem when using wanvideowrapper 1.3.5.

And after using --disable-pinned-memory and replacing the WanVideo Blockswap with WanVideo Enhanced BlockSwap (an ancient old node), the generation speed was slow, I was able to generate a 1920×1024×61 frame video using the 5060ti 16G graphics card!
https://github.com/user-attachments/assets/a975c627-5d41-4102-afd2-9ce83073cfe8

This is truly super amazing!!

I don't know how this setting manages to prevent the VRAM from being completely filled at the very beginning of the sampling process. It was like strangely achieving the generation of wan2.2 video within the DRAM memory.After replacing the wan video blockswap node, I can generate infinitely without encountering the problem of memory overflow.

Although the entire process took more than 1.5 hours,I also know that it is because the graphics card cannot directly access the memory. if using wanvideo blockswap node, there would definitely be an error of OOM and it would not be able to be generated.

The partial log of successfully generating 1920×1024×61 frames is as follows:
models:14B_High / Low FP8_scaled_e4m3fn + PusaV1_High_rank98 / PusaV1_rank512 + lightx2v_T2V_rank256 + Fun14B_Low_HPS2.1

Device: cuda:0   | Memory: 43,228.75 MB
Device: cpu      | Memory: 7,695,946.54 MB
Using 400 LoRA weight patches for WanVideo model
TeaCache: Using cache device: cpu
Rope function: comfy
---------- Sampling start ----------
61 frames at 1024x1920 (Input sequence length: 122880) with 7 steps
Generated new RoPE frequencies
100%|█████████████████████████| 7/7 [34:46<00:00, 298.14s/it]
---------- Sampling end ------------
TeaCache skipped: 0 conditional steps: []
TeaCache skipped: 0 unconditional steps: []
[Sampling] Max allocated memory: max_memory=11.024 GB
[Sampling] Max reserved memory: max_reserved=13.531 GB

Only succeeded in this mode: sageattn_compiled + Torch compile(inductor+default) under Sage attention 2.2,High 7 steps+ Low 8 steps，
If in sageattn_3 + Torch compile(inductor+default) mode is OOM, I don't know why

---

Regarding Question Two：

---

My Sage Attention 2.2 uses the pre-compiled version of woctordho, namely SageAttention2.2.post4.The work is functioning well and there are no problems.

At first, I intended to install the pre-compiled Sage Attention 3.0 from Mengqin, but it failed. Even though the version numbers matched, it was unable to output the correct image.The generated video is black and contains no content at all.
Then I attempted to compile Sage Attention 3.0 on the Windows platform using MSVC. I still used the method mentioned by Mengqin in the pull request #323 of thu-ml/SageAttention. I manually modified the files successful compilation,The compilation process strictly adhered to numerous restrictions.Specified the paths for using MSVC and CUDA,In addition, it also includes:

chcp 437 >nul
set VSLANG=1033
set TORCH_CUDA_ARCH_LIST=12.0
set MAX_PATH=4096
set MAX_JOB=1
set PYTHONLONGPATHS=1
set DISTUTILS_USE_SDK=1
set CXXFLAGS=/std:c++17 /O2 /D_CRT_SECURE_NO_WARNINGS
set NVCC_FLAGS=-O3 --use_fast_math -Xcompiler "/std:c++17"

Final it was able to generate the images correctly on WanVideoWrapper.
Then, during the test, I encountered a similar issue to #321 of thu-ml/SageAttention, which was also one of my problems.
At this point, I noticed that when using the default WAN2.2 workflow in the official settings and setting the Patch Sage Attention KJ node to use the sageattn3_per_block_mean mode, the generated video images were quite good and at least acceptable. However, when choosing the sageattn_3 mode in wanvideowrapper, the quality was very poor.
In order to test the image quality, I asked Gemini V3 to explain the principle while simultaneously modifying the file custom_nodes\ComfyUI-WanVideoWrapper\wanvideo\modules\model.py to the full Sage Attention 3.0 mode throughout.

Since I didn't understand the code, I gave the code to Gemini V3 for review. The AI added the following code in custom_nodes\ComfyUI-WanVideoWrapper\wanvideo\modules\attention.py:

elif attention_mode == 'sageattn_3':

    use_per_block = transformer_options.get("sageattn3_per_block", True)
    q_input = q.transpose(1,2).contiguous()
    k_input = k.transpose(1,2).contiguous()
    v_input = v.transpose(1,2).contiguous() 

    out = sageattn_blackwell(
        q_input, k_input, v_input,
        is_causal=causal,
        per_block_mean=use_per_block

The image quality has been improved! At a resolution of at least 448×800, it is acceptable!
https://github.com/user-attachments/assets/e6487553-46d1-43ec-b634-121f39fa7df7
The video is the result obtained by running in the modified pure sageatten_3 + WanVideo Torch Compile Settings mode,448×800
but at higher resolutions such as 640×1152, it can still be clearly noticed. at this point I modified model.py again to restore it to its original state. By setting the first and last steps to use the SageAttention2.2 mode, the video returned to normal!
At this point, the generation speed of the full version Sage Attention 3.0 is approximately equivalent to the speed of Sage Attention 2.2 when torch compile is enabled.
This did not fully demonstrate the capabilities of Sage Attention 3.0 in terms of efficiency, therefore I would like to have Sage Attention 3.0 also support WanVideo Torch Compile Settings node.
This time, I directly fed the entire python_embeded\Lib\site-packages\sageattn3\api.py file to Gemini V3. Google's AI completely overhauled this file:

import torch
import triton
import triton.language as tl
import torch.nn.functional as F
from typing import Tuple, Optional
from torch.nn.functional import scaled_dot_product_attention as sdpa
import fp4attn_cuda
import fp4quant_cuda
from torch.library import custom_op
_SA3_TRACE_COUNT = 0
# --- 第一步：注册算子 (用于 torch.compile 识别黑盒) ---

@custom_op("sageattn3::scaled_fp4_quant", mutates_args={"packed_fp4", "fp8_scale"})
def scaled_fp4_quant_wrapper(x: torch.Tensor, packed_fp4: torch.Tensor, fp8_scale: torch.Tensor, is_bf16: int) -> None:
    fp4quant_cuda.scaled_fp4_quant(x, packed_fp4, fp8_scale, is_bf16)

@scaled_fp4_quant_wrapper.register_fake
def _(x, packed_fp4, fp8_scale, is_bf16):
    return None

@custom_op("sageattn3::scaled_fp4_quant_permute", mutates_args={"packed_fp4", "fp8_scale"})
def scaled_fp4_quant_permute_wrapper(x: torch.Tensor, packed_fp4: torch.Tensor, fp8_scale: torch.Tensor, is_bf16: int) -> None:
    fp4quant_cuda.scaled_fp4_quant_permute(x, packed_fp4, fp8_scale, is_bf16)

@scaled_fp4_quant_permute_wrapper.register_fake
def _(x, packed_fp4, fp8_scale, is_bf16):
    return None

@custom_op("sageattn3::scaled_fp4_quant_trans", mutates_args={"packed_fp4", "fp8_scale"})
def scaled_fp4_quant_trans_wrapper(x: torch.Tensor, packed_fp4: torch.Tensor, fp8_scale: torch.Tensor, is_bf16: int) -> None:
    fp4quant_cuda.scaled_fp4_quant_trans(x, packed_fp4, fp8_scale, is_bf16)

@scaled_fp4_quant_trans_wrapper.register_fake
def _(x, packed_fp4, fp8_scale, is_bf16):
    return None

@custom_op("sageattn3::fp4attn_fwd", mutates_args={})
def fp4attn_fwd_wrapper(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, q_scale: torch.Tensor, k_scale: torch.Tensor, v_scale: torch.Tensor, 
    delta_s: torch.Tensor, KL: int, attn_mask: Optional[torch.Tensor], softmax_scale: float, is_causal: bool, 
    per_block_mean: bool, is_bf16: bool
) -> torch.Tensor:
    global _SA3_TRACE_COUNT
    _SA3_TRACE_COUNT += 1
    
    # 验证逻辑：每隔 100 次调用打印一次证据
    if _SA3_TRACE_COUNT % 100 == 0:
        # 获取输入张量的真实数据类型和量化状态作为证据
        print(f"\n[SA3.0-VERIFY] 正在调用底层 FP4 算子 | 总计调用次数: {_SA3_TRACE_COUNT}")
        print(f" - Q-Packed 形状: {q.shape} | Q-Scale 类型: {q_scale.dtype}")
        print(f" - Per-Block-Mean 补偿状态: {per_block_mean}")
        
    return fp4attn_cuda.fwd(q, k, v, q_scale, k_scale, v_scale, delta_s, KL, attn_mask, softmax_scale, is_causal, per_block_mean, is_bf16)[0]

@fp4attn_fwd_wrapper.register_fake
def _(q, k, v, q_scale, k_scale, v_scale, delta_s, KL, attn_mask, softmax_scale, is_causal, per_block_mean, is_bf16):
    # 强制指定输出形状为原始 HeadDim 128
    B, H, L, _ = q.shape
    return torch.empty((B, H, L, 128), device=q.device, dtype=torch.bfloat16 if is_bf16 else torch.float16)

# --- 第二步：逻辑实现 ---

@triton.jit
def group_mean_kernel(
    q_ptr, q_out_ptr, qm_out_ptr, B, H, L, D: tl.constexpr,
    stride_qb, stride_qh, stride_ql, stride_qd,
    stride_qmb, stride_qmh, stride_qml, stride_qmd,
    GROUP_SIZE: tl.constexpr
):
    pid_b, pid_h, pid_group = tl.program_id(0), tl.program_id(1), tl.program_id(2)
    offsets = pid_group * GROUP_SIZE + tl.arange(0, GROUP_SIZE)
    q_offsets = pid_b * stride_qb + pid_h * stride_qh + offsets[:, None] * stride_ql + tl.arange(0, D)[None, :] * stride_qd
    q_group = tl.load(q_ptr + q_offsets)
    qm_group = tl.sum(q_group, axis=0) / GROUP_SIZE
    tl.store(q_out_ptr + q_offsets, q_group - qm_group)
    qm_offset = pid_b * stride_qmb + pid_h * stride_qmh + pid_group * stride_qml + tl.arange(0, D) * stride_qmd
    tl.store(qm_out_ptr + qm_offset, qm_group)

def triton_group_mean(q: torch.Tensor):
    B, H, L, D = q.shape
    GROUP_SIZE = 128
    num_groups = L // GROUP_SIZE
    q_out, qm = torch.empty_like(q), torch.empty(B, H, num_groups, D, device=q.device, dtype=q.dtype)
    group_mean_kernel[(B, H, num_groups)](
        q, q_out, qm, B, H, L, D,
        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
        qm.stride(0), qm.stride(1), qm.stride(2), qm.stride(3),
        GROUP_SIZE=GROUP_SIZE
    )
    return q_out, qm

def preprocess_qkv(q, k, v, per_block_mean=True):
    def pad_128(x):
        L = x.size(2)
        pad = (128 - L % 128) % 128
        return F.pad(x, (0, 0, 0, pad)).contiguous() if pad > 0 else x.contiguous()
    
    k = k - k.mean(dim=-2, keepdim=True)
    q, k, v = map(pad_128, [q, k, v])
    if per_block_mean:
        q, qm = triton_group_mean(q)
    else:
        qm = q.mean(dim=-2, keepdim=True)
        q = q - qm
    delta_s = torch.matmul(qm, k.transpose(-2, -1)).to(torch.float32).contiguous()
    return q, k, v, delta_s

def scale_and_quant_fp4_all(x, mode="base"):
    B, H, N, D = x.shape
    # 根据模式决定 packed 形状
    p_shape = (B, H, D, N // 2) if mode == "trans" else (B, H, N, D // 2)
    s_shape = (B, H, D, N // 16) if mode == "trans" else (B, H, N, D // 16)
    packed = torch.empty(p_shape, device=x.device, dtype=torch.uint8)
    scale = torch.empty(s_shape, device=x.device, dtype=torch.float8_e4m3fn)
    
    if mode == "base": torch.ops.sageattn3.scaled_fp4_quant(x, packed, scale, 1)
    elif mode == "permute": torch.ops.sageattn3.scaled_fp4_quant_permute(x, packed, scale, 1)
    elif mode == "trans": torch.ops.sageattn3.scaled_fp4_quant_trans(x, packed, scale, 1)
    return packed, scale

def sageattn3_blackwell(q, k, v, attn_mask=None, is_causal=False, per_block_mean=True, **kwargs):
    if q.size(-1) >= 256: return sdpa(q, k, v, is_causal=is_causal)
    QL, KL, is_bf16 = q.size(2), k.size(2), q.dtype == torch.bfloat16
    
    q, k, v, delta_s = preprocess_qkv(q, k, v, per_block_mean)
    q_pk, q_sc = scale_and_quant_fp4_all(q, "base")
    k_pk, k_sc = scale_and_quant_fp4_all(k, "permute")
    v_pk, v_sc = scale_and_quant_fp4_all(v, "trans")
    
    o_fp4 = torch.ops.sageattn3.fp4attn_fwd(
        q_pk, k_pk, v_pk, q_sc, k_sc, v_sc, delta_s, KL, None, 128**-0.5, is_causal, per_block_mean, is_bf16
    )
    return o_fp4[:, :, :QL, :].contiguous()

!!!Please note that I have no idea what potential problems this change might have!!!
During this modification process, many revisions were made. I was only responsible for pasting and doing a little bit of analysis. I simply couldn't understand it anymore. Anyway, this modification successfully enabled Sage Attention 3.0 to start torch compile in the default mode. Moreover, I also had Gemini V3 incorporate verification code to confirm that this API was called.
Now, after I clear all the cache, when I first run Sage Attention 3.0 and Torch Compile for wan video sampling, a compilation process like the one shown in the following screenshot will occur.

(Don't worry about the sampling speed in the first step. The model library I placed in the HDD)
One minor issue is that if the resolution is modified, sometimes there will be a "torch compile" error. Restarting ComfyUI or clearing the cache can restore normal operation.
All the compilation processes will fail to compile if using the default tcc.exe.
In the startup configuration, I have explicitly specified MSVC for compilingcall "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat". Additionally, it is necessary to explicitly specify chcp 65001 and set PYTHONUTF8=1 in order to successfully enable it.

Until the end, I tried to have the AI help me activate reduce-overhead but it didn't work.
When attempting to run torch compile using the reduce-overhead mode, an error will occur:

RuntimeError: cudaMallocAsync does not yet support checkPoolLiveAllocations.

Even though I enabled the --disable-cuda-malloc command, it did not take effect (in the independent test script).The error will continue to occur:

Failed! Error: These storage data ptrs are not allocated in pool (0, 1) but should be {124591800320}

The response provided by Gemini v3 is: You are actually instructing PyTorch to bypass its own "memory pool" (Pool) and directly apply for addresses from the graphics card.
The conflict point: The reduce-overhead mode of torch.compile is an extremely "obsessive-compulsive" management mechanism. When recording CUDA Graphs, it defaults that all data must be in the monitored GPU memory pool (Pool 0 or Pool 1).
The error message indicates that {124591800320} is a specific physical address of the video memory. The compiler is saying: "I found this data at a certain location on the graphics card, but it is not in the 0th or 1st pool I preset. This does not comply with the recording specification, so I refuse to execute."
Path dependency error: On the Windows platform, the reduce-overhead mode of PyTorch is mutually exclusive with PYTORCH_NO_CUDA_MEMORY_CACHING=1. You cannot fix it by disabling the video memory pool because the recorder itself depends on the video memory pool.
Operator behavior is normal: The error occurs during the "recording phase" rather than the "execution phase", which indicates that your SA3.0 C++ operator itself is fine. The problem lies in how the PyTorch compiler manages the video memory addresses of these operators.
Driver/Environment Bottleneck: This error once again points to the incomplete support for cudaMallocAsync by the Windows driver.

At this point, I'm finally stuck at this step.

I feel that the 50xx series graphics cards are getting very close to being fully compatible with the Wanvideowrapper and Sage Attention 3.0 at full speed.I really like the wanvideowrapper node group and I'm very grateful to Kijia!