LTX2 Memory Efficient Sage Attention Patch crashes on RTX 5090 because SageAttention ops are registered under torch.ops, not _qattn_sm89 attributes

[Setmaster]

Describe the bug

On Windows with an RTX 5090, enabling LTX2MemoryEfficientSageAttentionPatch for an LTX2 workflow crashes at the first sampling step with:

AttributeError: module 'sageattention._qattn_sm89' has no attribute 'qk_int8_sv_f8_accum_f32_fuse_v_scale_attn'

The workflow runs when this node is bypassed. Native ComfyUI --use-sage-attention is also present in my launch args and is not the failing path here. The failure only starts when LTX2MemoryEfficientSageAttentionPatch replaces the LTX2 attn1.forward calls.

Why this looks like a KJNodes compatibility issue

In nodes/ltxv_nodes.py, the sm120 branch calls private SageAttention extension module attributes:

_qattn_sm89.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(...)

On this install, those functions are not exposed as attributes on sageattention._qattn_sm89, but the corresponding kernels are available as registered PyTorch custom ops after importing sageattention.core:

torch.ops.sageattention_qattn_sm89.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn
torch.ops.sageattention_qattn_sm89.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf
torch.ops.sageattention_qattn_sm89.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf

SageAttention itself is not generally broken on this machine. Both local smoke tests below succeed on the RTX 5090:

from sageattention import sageattn
from sageattention.core import sageattn_qk_int8_pv_fp8_cuda

o = sageattn(q, k, v, tensor_layout="NHD", is_causal=False)
o = sageattn_qk_int8_pv_fp8_cuda(
    q, k, v,
    tensor_layout="NHD",
    is_causal=False,
    qk_quant_gran="per_warp",
    pv_accum_dtype="fp32+fp16",
)

Environment

OS: Windows
Python: 3.13.7
GPU: NVIDIA GeForce RTX 5090
CUDA capability: (12, 0), SageAttention reports sm120
torch: 2.9.1+cu130
torch CUDA: 13.0
sageattention: 2.2.0+cu130torch2.9.0.post3
triton-windows: 3.5.1.post24
comfyui-kjnodes: 1.4.7 from Comfy Registry
ComfyUI: recent master when tested

Reproduction

1. Use an LTX2 / LTX2.3 workflow that includes LTX2MemoryEfficientSageAttentionPatch.
2. Enable the node.
3. Queue the prompt on an RTX 5090 / sm120 system with the environment above.
4. Sampling fails near step 0 with the _qattn_sm89 missing attribute error.

In my workflow, the node is LTX2MemoryEfficientSageAttentionPatch, with triton_kernels=True. Turning triton_kernels off is unlikely to address this specific error because that widget controls fused RoPE/helper kernels around the LTX2 forward path, not the failing SageAttention attention op call.

Relevant stack trace excerpt

File "<ComfyUI>\custom_nodes\comfyui-kjnodes\nodes\ltxv_nodes.py", line 1863, in ltx2_sageattn_forward
    o = _sageattn_int8_fp8_nhd(qkv, dtype)

File "<ComfyUI>\custom_nodes\comfyui-kjnodes\nodes\ltxv_nodes.py", line 1808, in _sageattn_int8_fp8_nhd
    _qattn_sm89.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(...)

AttributeError: module 'sageattention._qattn_sm89' has no attribute 'qk_int8_sv_f8_accum_f32_fuse_v_scale_attn'

Local verification

The direct module attribute is absent:

import sageattention.core as c
hasattr(c._qattn_sm89, "qk_int8_sv_f8_accum_f32_fuse_v_scale_attn")
# False

The registered op is present:

import torch
import sageattention.core
hasattr(torch.ops.sageattention_qattn_sm89, "qk_int8_sv_f8_accum_f32_fuse_v_scale_attn")
# True

The public local SageAttention function path also works:

public_sageattn (1, 128, 8, 64) torch.float16 True
public_fp8_cuda (1, 128, 8, 64) torch.float16 True

Suggested fix

This is intentionally narrower than PR #518. I am not asking to replace the LTX2 low-memory implementation with only SageAttention's high-level local Python function, since that was rejected for losing the memory benefit of this node.

The narrow fix would be to keep the existing _sageattn_int8_fp8_nhd low-memory path, but call the registered local PyTorch custom ops when the direct _qattn_sm89 attributes are absent. For example, resolve the op from either:

_qattn_sm89.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn

or:

torch.ops.sageattention_qattn_sm89.qk_int8_sv_f8_accum_f32_fuse_v_scale_attn

and update sageplus_sm89_available detection to check the same call surface. That should preserve the node's memory-conscious q/k/v lifetime behavior while supporting current Windows SageAttention wheels that register kernels via torch.ops.

Related items I found

• LTX2 Memory Efficient Sage Attention Patch fix (High level / Legacy API) #518 - related but broader high-level API / legacy fallback PR, closed unmerged
• [Bug] SageAttention performance regression - missing qk_quant_gran="per_warp" #325 - related sm120 / qk_quant_gran="per_warp" discussion, closed
• Patch Sage Attention KJ crashes with Fatal Python error: Aborted on Qwen Image Edit and Flux 2 Klein — native --use-sage-attention works fine on same workflows #670 - related RTX 5090 SageAttention KJ crash, but for the generic patch node and a hard abort, not this LTX2 _qattn_sm89 AttributeError