226 lines
10 KiB
Python
226 lines
10 KiB
Python
"""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")
|
|
DEFAULT_ATTENTION_BACKEND = os.getenv("H3_DEFAULT_ATTENTION", "sol_attn")
|
|
|
|
|
|
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_sol_attention_bshd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, is_causal: bool) -> torch.Tensor:
|
|
"""Run Sol-Attn on `[batch, sequence, heads, dim]` tensors and return the same layout."""
|
|
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[1] < min_tokens:
|
|
raise ValueError(f"{q.shape[1]} tokens < H3_SOL_MIN_TOKENS={min_tokens}")
|
|
return sol_attn(
|
|
q.contiguous(),
|
|
k.contiguous(),
|
|
v.contiguous(),
|
|
tau=tau,
|
|
thresh_type=thresh_type,
|
|
int8_qk=int8_qk,
|
|
int8_pv=int8_pv,
|
|
)
|
|
except Exception:
|
|
if os.getenv("H3_SOL_STRICT", "").lower() in {"1", "true", "yes", "on"}:
|
|
raise
|
|
fallback = os.getenv("H3_SOL_FALLBACK", "sage2")
|
|
hnd = run_attention(q.transpose(1, 2).contiguous(), k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous(), backend=fallback, is_causal=is_causal)
|
|
return hnd.transpose(1, 2)
|
|
|
|
|
|
def qkv_to_bshd(qkv: torch.Tensor, heads: int, head_dim: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""Split `[S, 3*H*D]` QKV into contiguous Sol-native `[1,S,H,D]` tensors."""
|
|
if os.getenv("H3_SOL_QKV_LAYOUT", "native").lower() != "fused":
|
|
raise RuntimeError("fused QKV layout disabled")
|
|
try:
|
|
from .nvfp4_quant import _vortex_scale_extension
|
|
|
|
return tuple(_vortex_scale_extension().qkv_to_bshd(qkv, heads, head_dim))
|
|
except Exception:
|
|
if os.getenv("H3_SOL_QKV_LAYOUT_STRICT", "").lower() in {"1", "true", "yes", "on"}:
|
|
raise
|
|
inner = heads * head_dim
|
|
sequence = qkv.shape[0]
|
|
q, k, v = qkv.split(inner, dim=-1)
|
|
return (
|
|
q.view(1, sequence, heads, head_dim).contiguous(),
|
|
k.view(1, sequence, heads, head_dim).contiguous(),
|
|
v.view(1, sequence, heads, head_dim).contiguous(),
|
|
)
|
|
|
|
|
|
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":
|
|
return run_sol_attention_bshd(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=is_causal).transpose(1, 2)
|
|
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 = DEFAULT_ATTENTION_BACKEND,
|
|
):
|
|
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 = DEFAULT_ATTENTION_BACKEND):
|
|
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)
|
|
if self.backend == "sol_attn":
|
|
if os.getenv("H3_SOL_QKV_LAYOUT", "native").lower() == "fused":
|
|
q, k, v = qkv_to_bshd(qkv, self.heads, self.head_dim)
|
|
else:
|
|
q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
|
|
out = run_sol_attention_bshd(q, k, v, is_causal=False)
|
|
return self.out_proj(out.reshape(sequence, inner).contiguous())
|
|
|
|
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())
|