Fix the SM120 Sage and ConvRot GELU kernels on Windows

lightx2v/common/ops/attn/sage_sm120_cached.py

@@ -8,11 +8,21 @@
 from .template import AttnWeightTemplate
 
 
+def _sage_kernels():
+    """SageAttention's fused and SM89 kernels; builds against the stable ABI (the Windows wheels) register them as torch ops."""
+    from sageattention import _fused, _qattn_sm89
+
+    if not hasattr(_fused, "quant_per_warp_int8_cuda"):
+        return torch.ops.sageattention_fused, torch.ops.sageattention_qattn_sm89
+    return _fused, _qattn_sm89
+
+
 @torch.library.custom_op("lightx2v::sage_sm120_transposed_v", mutates_args=(), device_types="cuda")
 def sage_sm120_transposed_v(q: torch.Tensor, k: torch.Tensor, transposed_v: torch.Tensor) -> torch.Tensor:
-    from sageattention import _fused, _qattn_sm89
     from sageattention.quant import per_warp_int8
 
+    _fused, _qattn_sm89 = _sage_kernels()
+
     dim = q.shape[-1]
     if dim not in (64, 128):
         raise ValueError("sage_attn2_sm120_cached requires head_dim 64 or 128")
@@ -32,9 +42,10 @@ def sage_sm120_transposed_v_fake(q, k, transposed_v):
 
 def prepare_sage_sm120_kv(k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """Quantize static [tokens, heads, dim] K/V exactly as native SM120 Sage."""
-    from sageattention import _fused
     from sageattention.quant import per_channel_fp8
 
+    _fused, _ = _sage_kernels()
+
     if k.ndim != 3 or k.shape != v.shape or k.shape[-1] not in (64, 128):
         raise ValueError("SM120 prepared K/V require matching [tokens, heads, dim] tensors with head_dim 64 or 128")
     if not k.is_cuda or k.device != v.device or k.dtype != v.dtype or k.dtype not in (torch.float16, torch.bfloat16):
@@ -53,8 +64,7 @@ def prepare_sage_sm120_kv(k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tenso
 @torch.library.custom_op("lightx2v::sage_sm120_cached_kv", mutates_args=(), device_types="cuda")
 def sage_sm120_cached_kv(q: torch.Tensor, ki: torch.Tensor, ks: torch.Tensor, vf: torch.Tensor, vs: torch.Tensor) -> torch.Tensor:
     """Attend to prepared static K/V, quantizing only the current query."""
-    from sageattention import _fused, _qattn_sm89
-
+    _fused, _qattn_sm89 = _sage_kernels()
     q = q.contiguous().unsqueeze(0)
     dim = q.shape[-1]
     if q.ndim != 4 or dim not in (64, 128):

lightx2v/common/ops/mm/fused_gelu_linear.py

@@ -130,7 +130,7 @@ def prepare_convrot_gelu_fusion(device=None):
                     _extension.lightx2v_prepare_convrot_gelu.restype = ctypes.c_int
                     _extension.lightx2v_gelu_convrot_quant.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p]
                     _extension.lightx2v_gelu_convrot_quant.restype = ctypes.c_int
-                except (OSError, RuntimeError) as error:
+                except (OSError, RuntimeError, AttributeError) as error:
                     _extension_error = error
                     return False
             if index not in _prepared_devices:

lightx2v_kernel/nvidia/csrc/quantization/convrot_gelu_quant.cu

@@ -161,11 +161,18 @@ __global__ void gelu_convrot_quant_kernel(const __nv_bfloat16* input, int8_t* qu
 
 } // namespace
 
-extern "C" int lightx2v_prepare_convrot_gelu() {
+// Loaded with ctypes: a Windows DLL exports only the functions marked for export.
+#if defined(_WIN32)
+#define LIGHTX2V_EXPORT __declspec(dllexport)
+#else
+#define LIGHTX2V_EXPORT
+#endif
+
+extern "C" LIGHTX2V_EXPORT int lightx2v_prepare_convrot_gelu() {
     return static_cast<int>(cudaFuncSetAttribute(gelu_convrot_quant_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSharedBytes));
 }
 
-extern "C" int lightx2v_gelu_convrot_quant(const void* input, void* quantized, void* scales, int rows, void* stream) {
+extern "C" LIGHTX2V_EXPORT int lightx2v_gelu_convrot_quant(const void* input, void* quantized, void* scales, int rows, void* stream) {
     if (rows == 0) return 0;
     gelu_convrot_quant_kernel<<<rows, kThreads, kSharedBytes, static_cast<cudaStream_t>(stream)>>>(
         static_cast<const __nv_bfloat16*>(input), static_cast<int8_t*>(quantized), static_cast<float*>(scales));