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

[cyberbobjr]

Describe the bug

Using the Patch Sage Attention KJ node causes an immediate hard crash (Fatal Python error: Aborted, exception code 0x80000003) at step 0/20 on the first attention call, with both Qwen Image Edit and Flux 2 Klein 9B workflows. All sage_attention modes fail:

• auto → crash
• sageattn_qk_int8_pv_fp16_cuda → crash
• sageattn_qk_int8_pv_fp16_triton → black image (no crash)

The same workflows run perfectly with ComfyUI's native --use-sage-attention flag (node bypassed), which suggests the issue is in the node's wrapper (model_optimization_nodes.py, sage_func line 36 / attention_sage line 84) rather than in sageattention itself.

Environment

• GPU: RTX 5090 (Blackwell, sm_120)
• OS: Windows 11, ComfyUI portable, Python 3.12 (embedded)
• torch: 2.12.0+cu130 (stable), torchvision 0.27.0+cu130, torchaudio 2.11.0+cu130
• sageattention: 2.2.0+cu130torch2.9.0andhigher.post4 (woct0rdho wheel)
• triton-windows: latest (upgraded during debugging, crash unchanged)
• comfyui-kjnodes: latest (git pull before testing)
• ComfyUI: latest
• Launch flags: --windows-standalone-build only (no --fast, no compile node, allow_compile: false on the patch node)

Ruled out during debugging

SageAttention itself works standalone in this exact environment. All of these pass without error:

# basic fp16
sageattn(torch.randn(1,8,1024,64, device='cuda', dtype=torch.float16))

# Qwen/Flux-like: bf16, head_dim 128, long seq
sageattn(torch.randn(1,24,4096,128, device='cuda', dtype=torch.bfloat16))

# non-aligned seq length
sageattn(torch.randn(1,24,7350,128, device='cuda', dtype=torch.bfloat16))

# cross q/k lengths
sageattn(q=(1,24,9216,128), k=v=(1,24,7350,128))

# NHD layout, non-contiguous (chunked from fused QKV), batch=2
sageattn(*torch.randn(1,4608,72,128, dtype=torch.bfloat16).chunk(3, dim=2), tensor_layout='NHD')
sageattn(torch.randn(2,4608,24,128, dtype=torch.bfloat16), tensor_layout='NHD')

Also ruled out:

• triton cache (cleared)
• torch/torchvision mismatch (clean stable trio reinstall)
• xformers (not installed)
• --fast flag (removed)
• torch.compile (no compile node in workflow, allow_compile: false)

To reproduce

1. Flux 2 Klein 9B fp8 text-to-image workflow (UNETLoader + CLIPLoader flux2 + SamplerCustomAdvanced, cfg 3, euler, Flux2Scheduler 20 steps, 1024×1024) — no ReferenceLatent nodes needed, plain t2i crashes.
2. Insert Patch Sage Attention KJ between model loader and sampler, mode sageattn_qk_int8_pv_fp16_cuda, allow_compile: false.
3. Queue prompt → crash at 0/20 on first attention call.

Same result with Qwen Image Edit (TextEncodeQwenImageEditPlus, Qwen Rapid AIO checkpoint).

Stack trace (Python side, identical in all crashes)

File ".../comfyui-kjnodes/nodes/model_optimization_nodes.py", line 36 in sage_func
File ".../torch/_dynamo/eval_frame.py", line 1298 in _fn
File ".../comfyui-kjnodes/nodes/model_optimization_nodes.py", line 84 in attention_sage
File ".../comfyui-kjnodes/nodes/model_optimization_nodes.py", line 126 in attention_override_sage
File ".../comfy/ldm/modules/attention.py", line 138 in wrapper
File ".../comfy/ldm/flux/math.py", line 14 in attention
...

Native crash side: abort raised from libtriton.pyd, then CUDA error surfaced via c10_cuda_check_implementation during tensor destruction.

Full stack trace (click to expand)

[INFO] Model Flux2 prepared for dynamic VRAM loading. 8658MB Staged. 0 patches attached. Force pre-loaded 80 weights: 49 KB.
  0%|                                                               | 0/20 [00:00<?, ?it/s,   Model Initializing ...  ]Fatal Python error: Aborted

Stack (most recent call first):
  File "F:\ComfyUI_windows_portable\ComfyUI\custom_nodes\comfyui-kjnodes\nodes\model_optimization_nodes.py", line 36 in sage_func
  File "F:\ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\_dynamo\eval_frame.py", line 1298 in _fn
  File "F:\ComfyUI_windows_portable\ComfyUI\custom_nodes\comfyui-kjnodes\nodes\model_optimization_nodes.py", line 84 in attention_sage
  File "F:\ComfyUI_windows_portable\ComfyUI\custom_nodes\comfyui-kjnodes\nodes\model_optimization_nodes.py", line 126 in attention_override_sage
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\ldm\modules\attention.py", line 138 in wrapper
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\ldm\flux\math.py", line 14 in attention
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\ldm\flux\layers.py", line 235 in forward
  File "F:\ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1789 in _call_impl
  File "F:\ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1778 in _wrapped_call_impl
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\ldm\flux\model.py", line 243 in forward_orig
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\ldm\flux\model.py", line 406 in _forward
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 113 in execute
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\ldm\flux\model.py", line 345 in forward
  File "F:\ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1789 in _call_impl
  File "F:\ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\nn\modules\module.py", line 1778 in _wrapped_call_impl
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\model_base.py", line 232 in _apply_model
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 113 in execute
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\model_base.py", line 188 in apply_model
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 334 in _calc_cond_batch
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 113 in execute
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 218 in _calc_cond_batch_outer
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 210 in calc_cond_batch
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 619 in sampling_function
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 1212 in predict_noise
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 113 in execute
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 1209 in outer_predict_noise
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 1202 in __call__
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 639 in __call__
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\k_diffusion\sampling.py", line 205 in sample_euler
  File "F:\ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\utils\_contextlib.py", line 124 in decorate_context
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 999 in sample
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 113 in execute
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 1229 in inner_sample
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 1254 in outer_sample
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\patcher_extension.py", line 113 in execute
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy\samplers.py", line 1316 in sample
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy_extras\nodes_custom_sampler.py", line 1048 in execute
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy_api\latest\_io.py", line 1892 in EXECUTE_NORMALIZED
  File "F:\ComfyUI_windows_portable\ComfyUI\comfy_api\internal\__init__.py", line 149 in wrapped_func
  File "F:\ComfyUI_windows_portable\ComfyUI\execution.py", line 299 in process_inputs
  File "F:\ComfyUI_windows_portable\ComfyUI\execution.py", line 311 in _async_map_node_over_list
  File "F:\ComfyUI_windows_portable\ComfyUI\execution.py", line 337 in get_output_data
  File "F:\ComfyUI_windows_portable\ComfyUI\execution.py", line 538 in execute
  File "F:\ComfyUI_windows_portable\ComfyUI\execution.py", line 780 in execute_async
  File "asyncio\events.py", line 88 in _run
  File "asyncio\base_events.py", line 1999 in _run_once
  File "asyncio\base_events.py", line 645 in run_forever
  File "asyncio\windows_events.py", line 322 in run_forever
  File "asyncio\base_events.py", line 678 in run_until_complete
  File "asyncio\runners.py", line 118 in run
  File "asyncio\runners.py", line 195 in run
  File "F:\ComfyUI_windows_portable\ComfyUI\execution.py", line 720 in execute
  File "F:\ComfyUI_windows_portable\ComfyUI\custom_nodes\comfyui-mnemic-nodes\utils\image_save_runtime_hook.py", line 14 in run
  File "F:\ComfyUI_windows_portable\ComfyUI\custom_nodes\comfyui_image_metadata_extension\modules\__init__.py", line 12 in run
  File "F:\ComfyUI_windows_portable\ComfyUI\main.py", line 340 in prompt_worker
  File "threading.py", line 1012 in run
  File "threading.py", line 1075 in _bootstrap_inner
  File "threading.py", line 1032 in _bootstrap

Extension modules: sqlalchemy.cyextension.collections, sqlalchemy.cyextension.immutabledict, sqlalchemy.cyextension.processors, sqlalchemy.cyextension.resultproxy, sqlalchemy.cyextension.util, greenlet._greenlet, markupsafe._speedups, yaml._yaml, PIL._imaging, multidict._multidict, yarl._quoting_c, propcache._helpers_c, _brotli, aiohttp._http_writer, aiohttp._http_parser, aiohttp._websocket.mask, aiohttp._websocket.reader_c, frozenlist._frozenlist, zstandard.backend_c, charset_normalizer.md, numpy.core._multiarray_umath, numpy.core._multiarray_tests, numpy.linalg._umath_linalg, numpy.fft._pocketfft_internal, numpy.random._common, numpy.random.bit_generator, numpy.random._bounded_integers, numpy.random._mt19937, numpy.random.mtrand, numpy.random._philox, numpy.random._pcg64, numpy.random._sfc64, numpy.random._generator, torch._C, torch._C._dynamo.autograd_compiler, torch._C._dynamo.eval_frame, torch._C._dynamo.guards, torch._C._dynamo.utils, torch._C._fft, torch._C._linalg, torch._C._nested, torch._C._nn, torch._C._sparse, torch._C._special, psutil._psutil_windows, PIL._imagingft, cython.cimports.libc.math, _cyutility, scipy._cyutility, scipy._lib._ccallback_c, scipy.ndimage._nd_image, scipy.ndimage._rank_filter_1d, scipy.special._ufuncs_cxx, scipy.special._ellip_harm_2, scipy.special._special_ufuncs, scipy.special._gufuncs, scipy.special._ufuncs, scipy.special._specfun, scipy.special._comb, _ni_label, scipy.ndimage._ni_label, regex._regex, sentencepiece._sentencepiece, scipy.integrate._odepack, scipy.integrate._quadpack, scipy.integrate._vode, scipy.integrate._dop, scipy.sparse._sparsetools, _csparsetools, scipy.sparse._csparsetools, scipy.linalg._fblas, scipy.linalg._flapack, scipy.linalg.cython_lapack, scipy.linalg._cythonized_array_utils, scipy.linalg._solve_toeplitz, scipy.linalg._batched_linalg, scipy.linalg._decomp_lu_cython, scipy.linalg._matfuncs_schur_sqrtm, scipy.linalg._matfuncs_expm, scipy.linalg._linalg_pythran, scipy.linalg.cython_blas, scipy.linalg._decomp_update, scipy.sparse.linalg._dsolve._superlu, scipy.sparse.linalg._eigen.arpack._arpacklib, scipy.sparse.linalg._propack, scipy.optimize._group_columns, scipy._lib.messagestream, scipy.optimize._trlib._trlib, scipy.optimize._lbfgsb, _moduleTNC, scipy.optimize._moduleTNC, scipy.optimize._slsqplib, scipy.optimize._minpack, scipy.optimize._lsq.givens_elimination, scipy.optimize._zeros, scipy._lib._uarray._uarray, scipy.linalg._decomp_interpolative, scipy.optimize._bglu_dense, scipy.optimize._lsap, scipy.spatial._ckdtree, scipy.spatial._qhull, scipy.spatial._voronoi, scipy.spatial._hausdorff, scipy.spatial._distance_wrap, scipy.spatial.transform._rotation_cy, scipy.spatial.transform._rigid_transform_cy, scipy.optimize._direct, scipy.interpolate._fitpack, scipy.interpolate._dfitpack, scipy.interpolate._dierckx, scipy.interpolate._ppoly, scipy.interpolate._interpnd, scipy.interpolate._rbfinterp_pythran, scipy.interpolate._rgi_cython, scipy.special.cython_special, scipy.stats._stats, scipy.stats._biasedurn, scipy.stats._stats_pythran, scipy.stats._levy_stable.levyst, scipy.stats._ansari_swilk_statistics, scipy.sparse.csgraph._tools, scipy.sparse.csgraph._shortest_path, scipy.sparse.csgraph._traversal, scipy.sparse.csgraph._min_spanning_tree, scipy.sparse.csgraph._flow, scipy.sparse.csgraph._matching, scipy.sparse.csgraph._reordering, scipy.stats._sobol, scipy.stats._qmc_cy, scipy.stats._rcont.rcont, scipy.stats._qmvnt_cy, av._core, av.logging, av.buffer, av.audio.format, av.error, av.dictionary, av.container.pyio, av.option, av.descriptor, av.format, av.index, av.utils, av.stream, av.container.streams, av.sidedata.encparams, av.sidedata.motionvectors, av.sidedata.sidedata, av.opaque, av.packet, av.container.input, av.container.output, av.container.core, av.codec.context, av.video.format, av.video.reformatter, av.plane, av.video.plane, av.video.frame, av.video.stream, av.codec.hwaccel, av.codec.codec, av.frame, av.audio.layout, av.audio.plane, av.audio.frame, av.audio.stream, av.filter.link, av.filter.context, av.filter.graph, av.filter.filter, av.filter.loudnorm, av.audio.resampler, av.audio.codeccontext, av.audio.fifo, av.bitstream, av.device, av.video.codeccontext, av.subtitles.stream, _cffi_backend, scipy.signal._sigtools, scipy.signal._max_len_seq_inner, scipy.signal._upfirdn_apply, scipy.signal._spline, scipy.signal._sosfilt, scipy.signal._peak_finding_utils, kiwisolver._cext, pandas._libs.tslibs.ccalendar, pandas._libs.tslibs.np_datetime, pandas._libs.tslibs.dtypes, pandas._libs.tslibs.base, pandas._libs.tslibs.nattype, pandas._libs.tslibs.timezones, pandas._libs.tslibs.fields, pandas._libs.tslibs.timedeltas, pandas._libs.tslibs.tzconversion, pandas._libs.tslibs.timestamps, pandas._libs.properties, pandas._libs.tslibs.offsets, pandas._libs.tslibs.strptime, pandas._libs.tslibs.parsing, pandas._libs.tslibs.conversion, pandas._libs.tslibs.period, pandas._libs.tslibs.vectorized, pandas._libs.ops_dispatch, pandas._libs.missing, pandas._libs.hashtable, pandas._libs.algos, pandas._libs.interval, pandas._libs.lib, pandas._libs.ops, pandas._libs.hashing, pandas._libs.arrays, pandas._libs.tslib, pandas._libs.sparse, pandas._libs.internals, pandas._libs.indexing, pandas._libs.index, pandas._libs.writers, pandas._libs.join, pandas._libs.window.aggregations, pandas._libs.window.indexers, pandas._libs.reshape, pandas._libs.groupby, pandas._libs.json, pandas._libs.parsers, pandas._libs.testing, sklearn.__check_build._check_build, lz4._version, lz4.frame._frame, sklearn._cyutility, sklearn.utils._isfinite, sklearn.utils.sparsefuncs_fast, sklearn.utils.murmurhash, sklearn.utils._openmp_helpers, sklearn.metrics.cluster._expected_mutual_info_fast, sklearn.preprocessing._csr_polynomial_expansion, sklearn.preprocessing._target_encoder_fast, sklearn.metrics._dist_metrics, sklearn.metrics._pairwise_distances_reduction._datasets_pair, sklearn.utils._cython_blas, sklearn.metrics._pairwise_distances_reduction._base, sklearn.metrics._pairwise_distances_reduction._middle_term_computer, sklearn.utils._heap, sklearn.utils._sorting, sklearn.metrics._pairwise_distances_reduction._argkmin, sklearn.metrics._pairwise_distances_reduction._argkmin_classmode, sklearn.utils._vector_sentinel, sklearn.metrics._pairwise_distances_reduction._radius_neighbors, sklearn.metrics._pairwise_distances_reduction._radius_neighbors_classmode, sklearn.metrics._pairwise_fast, skimage.measure._ccomp, skimage.measure._moments_cy, skimage.measure._find_contours_cy, skimage.measure._marching_cubes_lewiner_cy, websockets.speedups, lxml._elementpath, lxml.etree, pycocotools._mask, google._upb._message, grpc._cython.cygrpc, srsly.ujson.ujson, srsly.msgpack._epoch, srsly.msgpack._packer, srsly.msgpack._unpacker, h5py._errors, h5py.defs, h5py._objects, h5py.h5, h5py.utils, h5py.h5t, h5py.h5s, h5py.h5ac, h5py.h5p, h5py.h5r, h5py._proxy, h5py._conv, h5py.h5z, h5py.h5a, h5py.h5d, h5py.h5ds, h5py.h5g, h5py.h5i, h5py.h5o, h5py.h5f, h5py.h5fd, h5py.h5pl, h5py.h5l, h5py._selector, blis.cy, thinc.backends.cblas, cymem.cymem, preshed.maps, blis.py, thinc.backends.linalg, murmurhash.mrmr, thinc.backends.numpy_ops, thinc.layers.premap_ids, thinc.layers.sparselinear, spacy.symbols, preshed.bloom, spacy.strings, spacy.attrs, spacy.parts_of_speech, spacy.morphology, spacy.lexeme, spacy.tokens.morphanalysis, spacy.tokens.token, spacy.tokens.span, spacy.tokens.span_group, spacy.tokens._retokenize, spacy.tokens.doc, spacy.vectors, spacy.vocab, spacy.training.align, spacy.training.alignment_array, spacy.pipeline._parser_internals.nonproj, spacy.training.example, spacy.training.gold_io, spacy.matcher.levenshtein, spacy.matcher.matcher, spacy.matcher.dependencymatcher, spacy.matcher.phrasematcher, spacy.tokenizer, spacy.pipeline.pipe, spacy.pipeline.trainable_pipe, spacy.pipeline._parser_internals.stateclass, spacy.pipeline._parser_internals.transition_system, spacy.kb.kb, spacy.kb.candidate, spacy.kb.kb_in_memory, spacy.ml.parser_model, thinc.extra.search, spacy.pipeline._parser_internals._beam_utils, spacy.pipeline.transition_parser, spacy.pipeline._parser_internals.arc_eager, spacy.pipeline.dep_parser, spacy.pipeline._edit_tree_internals.edit_trees, spacy.pipeline.tagger, spacy.pipeline.morphologizer, spacy.pipeline._parser_internals.ner, spacy.pipeline.ner, spacy.pipeline.senter, spacy.pipeline.sentencizer, _win32sysloader, win32api, win32job, pymupdf._extra, pymupdf._mupdf, sklearn.utils._fast_dict, sklearn.cluster._hierarchical_fast, sklearn.cluster._k_means_common, sklearn.cluster._k_means_elkan, sklearn.cluster._k_means_lloyd, sklearn.cluster._k_means_minibatch, sklearn.neighbors._partition_nodes, sklearn.neighbors._ball_tree, sklearn.neighbors._kd_tree, sklearn.utils.arrayfuncs, sklearn.utils._random, sklearn.utils._seq_dataset, sklearn.linear_model._cd_fast, _loss, sklearn._loss._loss, sklearn.svm._liblinear, sklearn.svm._libsvm, sklearn.svm._libsvm_sparse, sklearn.linear_model._sag_fast, sklearn.utils._weight_vector, sklearn.linear_model._sgd_fast, sklearn.decomposition._online_lda_fast, sklearn.decomposition._cdnmf_fast, sklearn.cluster._dbscan_inner, sklearn.cluster._hdbscan._tree, sklearn.cluster._hdbscan._linkage, sklearn.cluster._hdbscan._reachability, sklearn._isotonic, sklearn.tree._utils, sklearn.tree._tree, sklearn.tree._partitioner, sklearn.tree._splitter, sklearn.tree._criterion, sklearn.neighbors._quad_tree, sklearn.manifold._barnes_hut_tsne, sklearn.manifold._utils, simsimd, stringzilla, pywt._extensions._dwt, pywt._extensions._cwt, pywt._extensions._pywt, pywt._extensions._swt, cuda_utils (total: 392)

Expected behavior

Either working sage attention (as with the native --use-sage-attention path on the same workflows/hardware), or a graceful fallback to SDPA instead of a hard abort.

Possibly related: #403 (sageattention3 crash on Blackwell with Qwen via the same node).

Thanks for the great nodes — happy to test any fix or provide more logs.

[kijai]

Can't replicate the issue, Flux2 klein 9b working fine, using latest wheel from https://github.com/woct0rdho/SageAttention:

Windows, 4090, torch 2.12+cu132:

1.77s/it to 1.41s/it

Linux, 5090, torch 2.12+cu130:

1.01 it/s to 1.30 it/s

It's more likely an issue with your system or sageattention install itself than anything to do with the node, --use-sage-attention not crashing doesn't necessarily mean it's working, does it give a speedup?

I won't be adding fallbacks as the purpose of these nodes is to enable it for models and situations where it does work.