"""Direct H3 self-attention using packed NVFP4 linears and SageAttention3.""" import os import torch import torch.nn.functional as functional from torch import nn from .checkpoint import H3Checkpoint from .nvfp4 import Nvfp4Linear AVAILABLE_BACKENDS = ("sage2", "sdpa", "sage3", "sage3_mean", "kj_sage_cuda", "kj_sage_triton", "kj_sage_fp8", "kj_sage_fp8pp", "kj_head_sliced", "sol_attn") PLANNED_BACKENDS = ("flash4", "easycache", "h3_cache", "kj_chunked_ffn") def attention_backend_status() -> dict[str, str]: """Report direct-runtime attention choices without importing ComfyUI nodes.""" status = {name: "available" for name in AVAILABLE_BACKENDS} status.update({"sol_attn": "experimental: sparse Triton attention for eligible non-causal H3 attention calls; falls back below H3_SOL_MIN_TOKENS unless H3_SOL_STRICT=1"}) status.update({"flash4": "planned: exact Blackwell kernel adapter"}) status.update({"easycache": "planned: approximate denoiser cache"}) status.update({"h3_cache": "planned: approximate H3-specific cache"}) status.update({"kj_chunked_ffn": "available: exact H3 MLP row chunking via H3_MLP_CHUNKS or runtime args"}) return status def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight.to(x.dtype), eps) def run_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, backend: str, is_causal: bool) -> torch.Tensor: """Run one `[batch, heads, sequence, dim]` attention operation.""" if backend == "sol_attn": try: from sol_kernel import sol_attn tau = float(os.getenv("H3_SOL_TAU", "1.3")) min_tokens = int(os.getenv("H3_SOL_MIN_TOKENS", "4096")) thresh_type = os.getenv("H3_SOL_THRESH_TYPE", "diag") int8_qk = os.getenv("H3_SOL_INT8_QK", "").lower() in {"1", "true", "yes", "on"} int8_pv = os.getenv("H3_SOL_INT8_PV", "").lower() in {"1", "true", "yes", "on"} if is_causal: raise ValueError("Sol-Attn backend only supports non-causal H3 attention") if q.shape[-1] != 128: raise ValueError(f"Sol-Attn requires head dim 128, got {q.shape[-1]}") if q.shape[2] < min_tokens: raise ValueError(f"{q.shape[2]} tokens < H3_SOL_MIN_TOKENS={min_tokens}") out = sol_attn( q.transpose(1, 2).contiguous(), k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous(), tau=tau, thresh_type=thresh_type, int8_qk=int8_qk, int8_pv=int8_pv, ) return out.transpose(1, 2) except Exception: if os.getenv("H3_SOL_STRICT", "").lower() in {"1", "true", "yes", "on"}: raise return run_attention(q, k, v, backend=os.getenv("H3_SOL_FALLBACK", "sage2"), is_causal=is_causal) if backend == "kj_head_sliced": head_slice_size = int(os.getenv("H3_HEAD_SLICE_SIZE", "8")) base_backend = os.getenv("H3_HEAD_SLICE_BACKEND", "sage2") if head_slice_size <= 0: raise ValueError("H3_HEAD_SLICE_SIZE must be positive") outputs = [ run_attention(q[:, start:start + head_slice_size], k[:, start:start + head_slice_size], v[:, start:start + head_slice_size], backend=base_backend, is_causal=is_causal) for start in range(0, q.shape[1], head_slice_size) ] return torch.cat(outputs, dim=1) if backend == "sage2": from sageattention import sageattn return sageattn(q, k, v, is_causal=is_causal, tensor_layout="HND", smooth_k=False) if backend == "sage3": from sageattn3 import sageattn3_blackwell return sageattn3_blackwell(q, k, v, is_causal=is_causal) if backend == "sage3_mean": from sageattn3 import sageattn3_blackwell return sageattn3_blackwell(q, k, v, is_causal=is_causal, per_block_mean=True) if backend == "kj_sage_cuda": from sageattention import sageattn_qk_int8_pv_fp16_cuda return sageattn_qk_int8_pv_fp16_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32", tensor_layout="HND") if backend == "kj_sage_triton": from sageattention import sageattn_qk_int8_pv_fp16_triton return sageattn_qk_int8_pv_fp16_triton(q, k, v, is_causal=is_causal, tensor_layout="HND") if backend == "kj_sage_fp8": from sageattention import sageattn_qk_int8_pv_fp8_cuda return sageattn_qk_int8_pv_fp8_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32+fp32", tensor_layout="HND") if backend == "kj_sage_fp8pp": from sageattention import sageattn_qk_int8_pv_fp8_cuda return sageattn_qk_int8_pv_fp8_cuda(q, k, v, is_causal=is_causal, pv_accum_dtype="fp32+fp16", tensor_layout="HND") if backend == "sdpa": return functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal) raise ValueError(f"Unsupported H3 attention backend: {backend}") def apply_split_half_rope(x: torch.Tensor, rotation: torch.Tensor) -> torch.Tensor: """Apply H3's split-half rotary table to `[batch, sequence, heads, dim]`.""" rotated_width = rotation.shape[-3] * 2 half = rotated_width // 2 if rotated_width > x.shape[-1]: raise ValueError("RoPE rotation width exceeds the attention head dimension.") pair = torch.stack((x[..., :half], x[..., half:rotated_width]), dim=-1) pair = torch.matmul(rotation.to(x.dtype), pair.unsqueeze(-1)).squeeze(-1) return torch.cat((pair[..., 0], pair[..., 1], x[..., rotated_width:]), dim=-1) def rms_rope_split_half_( q: torch.Tensor, k: torch.Tensor, rotation: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, eps: float, ) -> tuple[torch.Tensor, torch.Tensor]: """Run Comfy Kitchen's standalone fused H3 Q/K normalization and RoPE.""" import comfy_kitchen # Registers the independent CUDA extension operators. del comfy_kitchen torch.ops.comfy_kitchen.rms_rope_split_half_( q, k, rotation, q_weight, k_weight, eps, rotation.shape[-3] * 2 ) return q, k class H3SageAttention(nn.Module): """One MiniMax H3 attention module, independent of ComfyUI and Raylight.""" def __init__( self, qkv_proj: Nvfp4Linear, out_proj: Nvfp4Linear, q_norm_weight: torch.Tensor, k_norm_weight: torch.Tensor, *, heads: int = 56, head_dim: int = 128, eps: float = 1e-5, backend: str = "sage2", ): super().__init__() self.qkv_proj = qkv_proj self.out_proj = out_proj self.heads = heads self.head_dim = head_dim self.eps = eps if backend in PLANNED_BACKENDS: raise ValueError(f"H3 attention backend '{backend}' needs a standalone adapter and is not installed.") if backend not in AVAILABLE_BACKENDS: raise ValueError(f"Unsupported H3 attention backend: {backend}") self.backend = backend self.register_buffer("q_norm_weight", q_norm_weight, persistent=False) self.register_buffer("k_norm_weight", k_norm_weight, persistent=False) @classmethod def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str, *, output_dtype=torch.bfloat16, backend: str = "sage2"): return cls( checkpoint.nvfp4_linear(f"{prefix}.qkv_proj", output_dtype=output_dtype), checkpoint.nvfp4_linear(f"{prefix}.out_proj", output_dtype=output_dtype), checkpoint.tensor(f"{prefix}.q_norm.weight", dtype=output_dtype), checkpoint.tensor(f"{prefix}.k_norm.weight", dtype=output_dtype), backend=backend, ) def forward(self, x: torch.Tensor, rope_rotation: torch.Tensor) -> torch.Tensor: if x.ndim != 2: raise ValueError("H3 attention expects `[sequence, hidden]` input.") sequence = x.shape[0] inner = self.heads * self.head_dim q, k, v = self.qkv_proj(x).split(inner, dim=-1) q = q.view(1, sequence, self.heads, self.head_dim) k = k.view(1, sequence, self.heads, self.head_dim) v = v.view(1, sequence, self.heads, self.head_dim) q, k = rms_rope_split_half_(q, k, rope_rotation, self.q_norm_weight, self.k_norm_weight, self.eps) q = q.transpose(1, 2).contiguous() k = k.transpose(1, 2).contiguous() v = v.transpose(1, 2).contiguous() out = run_attention(q, k, v, backend=self.backend, is_causal=False) return self.out_proj(out.transpose(1, 2).reshape(sequence, inner).contiguous())