From bdee3f53551909a6258e4f47890deb69eb7c228e Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Sat, 5 Sep 2026 14:28:04 +0000 Subject: [PATCH 1/5] feat(musa): add Moorcat sparse attention backend --- .../common/ops/attn/dynamic_sparse_attn.py | 69 ++++- lightx2v/common/ops/attn/utils/sla_util.py | 260 +++++++++++++++++- lightx2v/utils/registry_factory.py | 2 + lightx2v_platform/ops/__init__.py | 3 + .../ops/attn/mthreads_musa/__init__.py | 1 + .../ops/attn/mthreads_musa/moorcat_sparse.py | 200 ++++++++++++++ lightx2v_platform/registry_factory.py | 1 + 7 files changed, 528 insertions(+), 8 deletions(-) create mode 100644 lightx2v_platform/ops/attn/mthreads_musa/__init__.py create mode 100644 lightx2v_platform/ops/attn/mthreads_musa/moorcat_sparse.py diff --git a/lightx2v/common/ops/attn/dynamic_sparse_attn.py b/lightx2v/common/ops/attn/dynamic_sparse_attn.py index eb9281e81..6a52015d2 100644 --- a/lightx2v/common/ops/attn/dynamic_sparse_attn.py +++ b/lightx2v/common/ops/attn/dynamic_sparse_attn.py @@ -1,12 +1,12 @@ import torch from loguru import logger -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, SPARSE_OPERATOR_REGISTER from .kernels.sla_kernel import _attention from .kernels.sla_kernel_ar import _attention_ar from .template import AttnWeightTemplate -from .utils.sla_util import get_block_map, get_cuda_arch +from .utils.sla_util import get_block_lut_nhd, get_block_lut_nhd_uncentered, get_block_map, get_cuda_arch from .utils.sla_util_blhd import get_block_map_blhd from .utils.sparge_util import block_map_incremental_lut_triton, block_map_ordinal_lut_triton, sage2_block_sparse_attn @@ -73,11 +73,18 @@ def __init__(self, config=None): self.sparsity_ratio = float(self.config.get("sparsity_ratio", type(self).sparsity_ratio)) self.operator = self.config.get("operator", type(self).operator) self.per_block_mean = bool(self.config.get("per_block_mean", type(self).per_block_mean)) + self.force_local_blocks = int(self.config.get("force_local_blocks", 0)) + self.operator_setting = dict(self.config.get("operator_setting", {})) + self.fixed_topk = self.config.get("topk") + if self.fixed_topk is not None: + self.fixed_topk = int(self.fixed_topk) + if self.fixed_topk <= 0: + raise ValueError(f"dynamic sparse attention topk must be positive, got {self.fixed_topk}") if not 0.0 <= self.sparsity_ratio < 1.0: raise ValueError(f"dynamic sparse attention sparsity_ratio must be in [0, 1), got {self.sparsity_ratio}") - self.arch = get_cuda_arch(torch.cuda.current_device()) + self.arch = None self.topk = 1 - self.sparsity_ratio if self.operator == "triton": self.BLKQ, self.BLKK = 64, 64 @@ -86,6 +93,7 @@ def __init__(self, config=None): self.BLKQ, self.BLKK = 128, 128 self.apply_func = self.apply_triton_ar elif self.operator == "sage2": + self.arch = get_cuda_arch(torch.cuda.current_device()) if self.arch == "sm90": self.BLKQ, self.BLKK = 64, 128 else: @@ -100,6 +108,14 @@ def __init__(self, config=None): elif self.operator == "magi": self.BLKQ, self.BLKK = 128, 128 self.apply_func = self.apply_magi + elif self.operator in SPARSE_OPERATOR_REGISTER: + if self.fixed_topk is not None: + self.operator_setting.setdefault("topk", self.fixed_topk) + self.sparse_operator = SPARSE_OPERATOR_REGISTER[self.operator](self.operator_setting) + self.fixed_topk = getattr(self.sparse_operator, "topk", self.fixed_topk) + self.BLKQ = self.sparse_operator.q_block_size + self.BLKK = self.sparse_operator.k_block_size + self.apply_func = self.apply_registered_operator else: raise NotImplementedError(f"Not supported SLA operator: {self.operator}.") @@ -188,6 +204,53 @@ def apply_sage2( out = out.transpose(1, 2).reshape(max_seqlen_q, -1) return out + def apply_registered_operator( + self, + q, + k, + v, + cu_seqlens_q=None, + cu_seqlens_kv=None, + max_seqlen_q=None, + max_seqlen_kv=None, + **kwargs, + ): + if getattr(self.sparse_operator, "block_indices_only", False): + get_lut = get_block_lut_nhd if getattr(self.sparse_operator, "center_k", True) else get_block_lut_nhd_uncentered + block_indices, _ = get_lut( + q, + k, + topk_ratio=self.topk, + BLKQ=self.BLKQ, + BLKK=self.BLKK, + topk=self.fixed_topk, + force_local_blocks=self.force_local_blocks, + ) + sparse_map = None + else: + q_for_mask = q.unsqueeze(0).transpose(1, 2).contiguous() + k_for_mask = k.unsqueeze(0).transpose(1, 2).contiguous() + sparse_map, block_indices, _ = get_block_map( + q_for_mask, + k_for_mask, + topk_ratio=self.topk, + BLKQ=self.BLKQ, + BLKK=self.BLKK, + topk=self.fixed_topk, + ) + return self.sparse_operator( + q, + k, + v, + sparse_map, + block_indices=block_indices, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + **kwargs, + ) + def apply_sage3( self, q, diff --git a/lightx2v/common/ops/attn/utils/sla_util.py b/lightx2v/common/ops/attn/utils/sla_util.py index d717d9a94..4472abde2 100755 --- a/lightx2v/common/ops/attn/utils/sla_util.py +++ b/lightx2v/common/ops/attn/utils/sla_util.py @@ -44,7 +44,7 @@ def mean_pool(x, BLK): return x_mean -def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64): +def get_block_lut(q, k, topk_ratio, BLKQ=64, BLKK=64, topk=None): arg_k = k - torch.mean(k, dim=-2, keepdim=True) # smooth-k technique in SageAttention pooled_qblocks = mean_pool(q, BLKQ) pooled_kblocks = mean_pool(arg_k, BLKK) @@ -57,14 +57,27 @@ def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64): repeat_factor = num_q_heads // num_kv_heads pooled_kblocks = pooled_kblocks.repeat_interleave(repeat_factor, dim=1) - pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2) + if topk is not None and int(topk) > 16: + pooled_score = pooled_qblocks.float() @ pooled_kblocks.float().transpose(-1, -2) + else: + pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2) - K = pooled_score.shape[-1] + k_blocks = pooled_score.shape[-1] # Match the training router: short sequences still retain one key block. - topk = max(1, min(K, int(topk_ratio * K))) + if topk is None: + topk = int(topk_ratio * k_blocks) + topk = max(1, min(k_blocks, int(topk))) lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices + return lut, topk - sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8) + +def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64, topk=None): + lut, topk = get_block_lut(q, k, topk_ratio, BLKQ, BLKK, topk) + sparse_map = torch.zeros( + (*lut.shape[:-1], (k.shape[-2] + BLKK - 1) // BLKK), + device=lut.device, + dtype=torch.int8, + ) sparse_map.scatter_(-1, lut, 1) return sparse_map, lut, topk @@ -72,3 +85,240 @@ def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64): def get_cuda_arch(device_index): major, minor = torch.cuda.get_device_capability(device_index) return f"sm{major}{minor}" + + +@triton.jit(do_not_specialize=("L",)) +def compress_nhd_kernel( + X, + XM, + L, + H: tl.constexpr, + D: tl.constexpr, + BLOCK_L: tl.constexpr, +): + idx_l = tl.program_id(0) + idx_h = tl.program_id(1) + + offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L) + offs_d = tl.arange(0, D) + x = tl.load( + X + offs_l[:, None] * H * D + idx_h * D + offs_d[None, :], + mask=offs_l[:, None] < L, + other=0.0, + ) + + nx = min(BLOCK_L, L - idx_l * BLOCK_L) + x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx + l_blocks = (L + BLOCK_L - 1) // BLOCK_L + tl.store( + XM + idx_h * l_blocks * D + idx_l * D + offs_d, + x_mean.to(XM.dtype.element_ty), + ) + + +@triton.jit(do_not_specialize=("L",)) +def compress_centered_nhd_kernel( + X, + X_MEAN, + XM, + L, + H: tl.constexpr, + D: tl.constexpr, + BLOCK_L: tl.constexpr, +): + idx_l = tl.program_id(0) + idx_h = tl.program_id(1) + + offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L) + offs_d = tl.arange(0, D) + valid = offs_l[:, None] < L + x = tl.load( + X + offs_l[:, None] * H * D + idx_h * D + offs_d[None, :], + mask=valid, + other=0.0, + ) + x_mean = tl.load(X_MEAN + idx_h * D + offs_d) + centered = (x - x_mean[None, :]).to(X.dtype.element_ty) + centered = tl.where(valid, centered, 0.0) + + nx = min(BLOCK_L, L - idx_l * BLOCK_L) + block_mean = tl.sum(centered, axis=0, dtype=tl.float32) / nx + l_blocks = (L + BLOCK_L - 1) // BLOCK_L + tl.store( + XM + idx_h * l_blocks * D + idx_l * D + offs_d, + block_mean.to(XM.dtype.element_ty), + ) + + +def mean_pool_nhd(x, BLK): + assert x.ndim == 3 and x.is_contiguous() + + L, H, D = x.shape + L_BLOCKS = (L + BLK - 1) // BLK + x_mean = torch.empty((1, H, L_BLOCKS, D), device=x.device, dtype=x.dtype) + + compress_nhd_kernel[(L_BLOCKS, H)](x, x_mean, L, H, D, BLK) + return x_mean + + +def mean_pool_centered_nhd(x, x_mean, BLK): + assert x.ndim == 3 and x.is_contiguous() + assert x_mean.shape == x.shape[1:] + + L, H, D = x.shape + L_BLOCKS = (L + BLK - 1) // BLK + block_mean = torch.empty((1, H, L_BLOCKS, D), device=x.device, dtype=x.dtype) + + compress_centered_nhd_kernel[(L_BLOCKS, H)]( + x, + x_mean, + block_mean, + L, + H, + D, + BLK, + ) + return block_mean + + +def _force_local_score_blocks(pooled_score, count): + count = min(int(count), int(pooled_score.shape[-1])) + if count <= 0: + return pooled_score + q_blocks = int(pooled_score.shape[-2]) + k_blocks = int(pooled_score.shape[-1]) + centers = (torch.arange(q_blocks, device=pooled_score.device) * 2 + 1) * k_blocks // (2 * q_blocks) + offsets = torch.arange(count, device=pooled_score.device) - count // 2 + local_indices = (centers[:, None] + offsets[None, :]).clamp(0, k_blocks - 1) + local_indices = local_indices.view(*((1,) * (pooled_score.ndim - 2)), q_blocks, count).expand(*pooled_score.shape[:-2], q_blocks, count) + return pooled_score.scatter( + -1, + local_indices, + torch.finfo(pooled_score.dtype).max, + ) + + +def _get_block_lut_nhd( + q, + k, + topk_ratio, + BLKQ=64, + BLKK=64, + topk=None, + force_local_blocks=0, +): + pooled_qblocks = mean_pool_nhd(q, BLKQ) + k_mean = torch.mean(k, dim=0) + pooled_kblocks = mean_pool_centered_nhd(k, k_mean, BLKK) + + num_q_heads = q.size(1) + num_kv_heads = k.size(1) + if num_q_heads != num_kv_heads: + assert num_q_heads % num_kv_heads == 0, f"Number of Q heads ({num_q_heads}) must be divisible by number of KV heads ({num_kv_heads})" + pooled_kblocks = pooled_kblocks.repeat_interleave( + num_q_heads // num_kv_heads, + dim=1, + ) + + if topk is not None and int(topk) > 16: + pooled_score = pooled_qblocks.float() @ pooled_kblocks.float().transpose(-1, -2) + else: + pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2) + pooled_score = _force_local_score_blocks(pooled_score, force_local_blocks) + k_blocks = pooled_score.shape[-1] + if topk is None: + topk = int(topk_ratio * k_blocks) + topk = max(1, min(k_blocks, int(topk))) + return torch.topk(pooled_score, topk, dim=-1, sorted=False).indices, topk + + +def get_block_lut_nhd_uncentered( + q, + k, + topk_ratio, + BLKQ=64, + BLKK=64, + topk=None, + force_local_blocks=0, +): + pooled_qblocks = mean_pool_nhd(q, BLKQ) + pooled_kblocks = mean_pool_nhd(k, BLKK) + + num_q_heads = q.size(1) + num_kv_heads = k.size(1) + if num_q_heads != num_kv_heads: + assert num_q_heads % num_kv_heads == 0, f"Number of Q heads ({num_q_heads}) must be divisible by number of KV heads ({num_kv_heads})" + pooled_kblocks = pooled_kblocks.repeat_interleave( + num_q_heads // num_kv_heads, + dim=1, + ) + + if topk is not None and int(topk) > 16: + pooled_score = pooled_qblocks.float() @ pooled_kblocks.float().transpose(-1, -2) + else: + pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2) + pooled_score = _force_local_score_blocks(pooled_score, force_local_blocks) + k_blocks = pooled_score.shape[-1] + if topk is None: + topk = int(topk_ratio * k_blocks) + topk = max(1, min(k_blocks, int(topk))) + return torch.topk(pooled_score, topk, dim=-1, sorted=False).indices, topk + + +@torch.library.custom_op( + "lightx2v::block_lut_nhd", + mutates_args=(), + device_types="cuda", +) +def block_lut_nhd( + q: torch.Tensor, + k: torch.Tensor, + topk_ratio: float, + block_q: int, + block_k: int, + topk: int, +) -> torch.Tensor: + return _get_block_lut_nhd(q, k, topk_ratio, block_q, block_k, topk)[0] + + +@block_lut_nhd.register_fake +def _block_lut_nhd_fake( + q: torch.Tensor, + k: torch.Tensor, + topk_ratio: float, + block_q: int, + block_k: int, + topk: int, +) -> torch.Tensor: + q_blocks = (q.shape[0] + block_q - 1) // block_q + return torch.empty( + (1, q.shape[1], q_blocks, topk), + dtype=torch.int64, + device=q.device, + ) + + +def get_block_lut_nhd( + q, + k, + topk_ratio, + BLKQ=64, + BLKK=64, + topk=None, + force_local_blocks=0, +): + if topk is not None and int(topk) > 16 and int(force_local_blocks) == 0: + topk = int(topk) + return ( + block_lut_nhd(q, k, float(topk_ratio), BLKQ, BLKK, topk), + topk, + ) + return _get_block_lut_nhd( + q, + k, + topk_ratio, + BLKQ, + BLKK, + topk, + force_local_blocks, + ) diff --git a/lightx2v/utils/registry_factory.py b/lightx2v/utils/registry_factory.py index 7f0d4895c..bab8f81f6 100755 --- a/lightx2v/utils/registry_factory.py +++ b/lightx2v/utils/registry_factory.py @@ -9,6 +9,7 @@ PLATFORM_MM_WEIGHT_REGISTER, PLATFORM_RMS_WEIGHT_REGISTER, PLATFORM_ROPE_REGISTER, + PLATFORM_SPARSE_OPERATOR_REGISTER, ) @@ -93,6 +94,7 @@ def merge(self, other_register): SPARSE_OPERATOR_REGISTER = Register() ATTN_WEIGHT_REGISTER.merge(PLATFORM_ATTN_WEIGHT_REGISTER) +SPARSE_OPERATOR_REGISTER.merge(PLATFORM_SPARSE_OPERATOR_REGISTER) A2A_BACKEND_REGISTER.merge(PLATFORM_A2A_BACKEND_REGISTER) COMPILE_BACKEND_REGISTER.merge(PLATFORM_COMPILE_BACKEND_REGISTER) MM_WEIGHT_REGISTER.merge(PLATFORM_MM_WEIGHT_REGISTER) diff --git a/lightx2v_platform/ops/__init__.py b/lightx2v_platform/ops/__init__.py index 175d4deb1..cb696cffd 100755 --- a/lightx2v_platform/ops/__init__.py +++ b/lightx2v_platform/ops/__init__.py @@ -51,4 +51,7 @@ from .norm.iluvatar_cuda import * from .rope.iluvatar_cuda import * elif PLATFORM == "musa": + # Register platform attention operators before the framework registries + # take their one-time snapshot. + from .attn.mthreads_musa import * from .mm.mthreads_musa import * diff --git a/lightx2v_platform/ops/attn/mthreads_musa/__init__.py b/lightx2v_platform/ops/attn/mthreads_musa/__init__.py new file mode 100644 index 000000000..f156bb282 --- /dev/null +++ b/lightx2v_platform/ops/attn/mthreads_musa/__init__.py @@ -0,0 +1 @@ +from .moorcat_sparse import MusaMoorcatSparseOperator diff --git a/lightx2v_platform/ops/attn/mthreads_musa/moorcat_sparse.py b/lightx2v_platform/ops/attn/mthreads_musa/moorcat_sparse.py new file mode 100644 index 000000000..3eb6dbc9d --- /dev/null +++ b/lightx2v_platform/ops/attn/mthreads_musa/moorcat_sparse.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +import math + +import torch +from loguru import logger + +from lightx2v_platform.registry_factory import PLATFORM_SPARSE_OPERATOR_REGISTER + +try: + from moorcat._C import blocksparse as _blocksparse +except (ImportError, OSError) as exc: + logger.info(f"Moorcat block-sparse attention is unavailable: {exc}") + _blocksparse = None + + +_BLOCK_SIZE = 128 +_RUNTIME_CACHE = {} +_VALIDATED_CONTRACTS = set() + + +def _runtime_tensors( + q: torch.Tensor, + seq_len: int, + heads: int, + q_blocks: int, + topk: int, +): + key = (q.device, seq_len, heads, q_blocks, topk) + cached = _RUNTIME_CACHE.get(key) + if cached is None: + cached = ( + torch.arange(1, heads + 1, dtype=torch.int32, device=q.device), + torch.full( + (1, heads, q_blocks), + topk, + dtype=torch.int32, + device=q.device, + ), + torch.tensor([0, seq_len], dtype=torch.int32, device=q.device), + ) + _RUNTIME_CACHE[key] = cached + return cached + + +def _validate_lut_once( + block_indices: torch.Tensor, + *, + heads: int, + q_blocks: int, + kv_blocks: int, + topk: int, +) -> None: + expected = (1, heads, q_blocks, topk) + if tuple(block_indices.shape) != expected: + raise ValueError(f"Moorcat Q128 LUT must have shape {expected}, got {tuple(block_indices.shape)}") + contract = (expected, kv_blocks, block_indices.device) + if contract in _VALIDATED_CONTRACTS: + return + min_index = int(block_indices.min().item()) + max_index = int(block_indices.max().item()) + if min_index < 0 or max_index >= kv_blocks: + raise ValueError(f"Moorcat Q128 LUT indices must be in [0, {kv_blocks}), got [{min_index}, {max_index}]") + ordered = torch.sort(block_indices, dim=-1).values + if not bool(torch.all(ordered[..., 1:] != ordered[..., :-1]).item()): + raise ValueError("Moorcat Q128 LUT rows must contain distinct K blocks") + _VALIDATED_CONTRACTS.add(contract) + + +@torch.library.custom_op( + "lightx2v::musa_moorcat_sparse", + mutates_args=(), + device_types="cuda", +) +def musa_moorcat_sparse( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_indices: torch.Tensor, + topk: int, + softmax_scale: float, + validate_lut: bool, +) -> torch.Tensor: + if _blocksparse is None: + raise RuntimeError("Moorcat is not importable. Install the Moorcat package supplied with the official MUSA SGL image.") + + seq_len, heads, _ = q.shape + q_blocks = math.ceil(seq_len / _BLOCK_SIZE) + kv_blocks = math.ceil(k.shape[0] / _BLOCK_SIZE) + if validate_lut: + _validate_lut_once( + block_indices, + heads=heads, + q_blocks=q_blocks, + kv_blocks=kv_blocks, + topk=topk, + ) + + indices = block_indices.to(dtype=torch.int32).contiguous() + head_mask_type, counts, cu_seqlens = _runtime_tensors(q, seq_len, heads, q_blocks, topk) + + padded_kv_len = kv_blocks * _BLOCK_SIZE + if padded_kv_len != k.shape[0]: + k, v, _ = _blocksparse._pad_kv_tail_storage_to_tile128(k, v, padded_kv_len) + + out, _ = _blocksparse.block_sparse_attn_fwd_indexed( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + head_mask_type=head_mask_type, + streaming_info=None, + q2k_block_sparse_index=indices, + q2k_block_sparse_num=counts, + max_kv_blocks_per_q=topk, + max_seqlen_q=seq_len, + max_seqlen_k=seq_len, + softmax_scale=softmax_scale, + is_causal=False, + block_size=_BLOCK_SIZE, + ) + return out + + +@musa_moorcat_sparse.register_fake +def _musa_moorcat_sparse_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_indices: torch.Tensor, + topk: int, + softmax_scale: float, + validate_lut: bool, +) -> torch.Tensor: + return torch.empty_like(q) + + +@PLATFORM_SPARSE_OPERATOR_REGISTER("musa_moorcat_sparse") +class MusaMoorcatSparseOperator: + """Run a Q128/K128 block map with Moorcat's BF16 sparse kernel.""" + + q_block_size = _BLOCK_SIZE + k_block_size = _BLOCK_SIZE + block_indices_only = True + + def __init__(self, operator_setting=None): + setting = dict(operator_setting or {}) + self.topk = int(setting.get("topk", 16)) + self.center_k = bool(setting.get("center_k", True)) + self.validate_lut = bool(setting.get("validate_lut", setting.get("validate_mask", True))) + if self.topk <= 0: + raise ValueError(f"Moorcat topk must be positive, got {self.topk}") + + @torch.compiler.disable + def __call__( + self, + q, + k, + v, + mask, + block_indices=None, + cu_seqlens_q=None, + cu_seqlens_kv=None, + max_seqlen_q=None, + max_seqlen_kv=None, + **kwargs, + ): + if _blocksparse is None: + raise RuntimeError("The official Moorcat binary is unavailable. Install the Moorcat package supplied with the official MUSA SGL image.") + if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: + raise ValueError("Moorcat requires flattened [tokens, heads, head_dim] Q/K/V") + if q.shape != k.shape or q.shape != v.shape: + raise ValueError(f"Moorcat supports self-attention only, got q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}") + if q.dtype != torch.bfloat16 or k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("Moorcat requires BF16 Q/K/V") + if q.shape[-1] != 128: + raise ValueError(f"Moorcat requires head_dim=128, got {q.shape[-1]}") + if bool(kwargs.get("causal", False)): + raise ValueError("Moorcat supports non-causal attention only") + if cu_seqlens_q is not None and cu_seqlens_q.numel() != 2: + raise ValueError("Moorcat supports one packed sequence only") + if cu_seqlens_kv is not None and cu_seqlens_kv.numel() != 2: + raise ValueError("Moorcat supports one packed sequence only") + if block_indices is None: + raise ValueError("Moorcat requires Q128/K128 block indices") + + softmax_scale = kwargs.get("softmax_scale") + if softmax_scale is None: + softmax_scale = q.shape[-1] ** -0.5 + out = musa_moorcat_sparse( + q.contiguous(), + k.contiguous(), + v.contiguous(), + block_indices, + self.topk, + float(softmax_scale), + self.validate_lut, + ) + return out.reshape(q.shape[0], -1) diff --git a/lightx2v_platform/registry_factory.py b/lightx2v_platform/registry_factory.py index 39e8599d8..098ede2b3 100755 --- a/lightx2v_platform/registry_factory.py +++ b/lightx2v_platform/registry_factory.py @@ -66,6 +66,7 @@ def merge(self, other_register): PLATFORM_DEVICE_REGISTER = Register() PLATFORM_ATTN_WEIGHT_REGISTER = Register() +PLATFORM_SPARSE_OPERATOR_REGISTER = Register() PLATFORM_MM_WEIGHT_REGISTER = Register() PLATFORM_RMS_WEIGHT_REGISTER = Register() PLATFORM_LAYERNORM_WEIGHT_REGISTER = Register() From d91843656780a7cb2aaaf37eb9736054f49e9122 Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Sun, 6 Sep 2026 09:21:39 +0000 Subject: [PATCH 2/5] refactor(musa): simplify sparse attention path --- .../common/ops/attn/dynamic_sparse_attn.py | 34 +++-- lightx2v/common/ops/attn/utils/sla_util.py | 120 ++++++------------ .../ops/attn/mthreads_musa/moorcat_sparse.py | 66 ++-------- 3 files changed, 68 insertions(+), 152 deletions(-) diff --git a/lightx2v/common/ops/attn/dynamic_sparse_attn.py b/lightx2v/common/ops/attn/dynamic_sparse_attn.py index 6a52015d2..1ae71964a 100644 --- a/lightx2v/common/ops/attn/dynamic_sparse_attn.py +++ b/lightx2v/common/ops/attn/dynamic_sparse_attn.py @@ -69,22 +69,14 @@ class DynamicSparseAttnWeight(AttnWeightTemplate): per_block_mean = False def __init__(self, config=None): - self.config = dict(config or {}) - self.sparsity_ratio = float(self.config.get("sparsity_ratio", type(self).sparsity_ratio)) - self.operator = self.config.get("operator", type(self).operator) - self.per_block_mean = bool(self.config.get("per_block_mean", type(self).per_block_mean)) - self.force_local_blocks = int(self.config.get("force_local_blocks", 0)) - self.operator_setting = dict(self.config.get("operator_setting", {})) - self.fixed_topk = self.config.get("topk") - if self.fixed_topk is not None: - self.fixed_topk = int(self.fixed_topk) - if self.fixed_topk <= 0: - raise ValueError(f"dynamic sparse attention topk must be positive, got {self.fixed_topk}") + config = config or {} + self.sparsity_ratio = config.get("sparsity_ratio", type(self).sparsity_ratio) + self.operator = config.get("operator", type(self).operator) + self.per_block_mean = config.get("per_block_mean", type(self).per_block_mean) if not 0.0 <= self.sparsity_ratio < 1.0: raise ValueError(f"dynamic sparse attention sparsity_ratio must be in [0, 1), got {self.sparsity_ratio}") - self.arch = None self.topk = 1 - self.sparsity_ratio if self.operator == "triton": self.BLKQ, self.BLKK = 64, 64 @@ -109,18 +101,22 @@ def __init__(self, config=None): self.BLKQ, self.BLKK = 128, 128 self.apply_func = self.apply_magi elif self.operator in SPARSE_OPERATOR_REGISTER: + self.fixed_topk = config.get("topk") + if self.fixed_topk is not None and self.fixed_topk <= 0: + raise ValueError(f"dynamic sparse attention topk must be positive, got {self.fixed_topk}") + + operator_setting = config.get("operator_setting", {}).copy() if self.fixed_topk is not None: - self.operator_setting.setdefault("topk", self.fixed_topk) - self.sparse_operator = SPARSE_OPERATOR_REGISTER[self.operator](self.operator_setting) + operator_setting.setdefault("topk", self.fixed_topk) + self.sparse_operator = SPARSE_OPERATOR_REGISTER[self.operator](operator_setting) self.fixed_topk = getattr(self.sparse_operator, "topk", self.fixed_topk) + self.force_local_blocks = config.get("force_local_blocks", 0) self.BLKQ = self.sparse_operator.q_block_size self.BLKK = self.sparse_operator.k_block_size self.apply_func = self.apply_registered_operator else: raise NotImplementedError(f"Not supported SLA operator: {self.operator}.") - # logger.info(f"DynamicSparseAttnWeight: sparsity_ratio={self.sparsity_ratio}, operator={self.operator}, topk={self.topk}, BLKQ={self.BLKQ}, BLKK={self.BLKK}") - def apply( self, q, @@ -267,7 +263,7 @@ def apply_sage3( k = k.unsqueeze(0).transpose(1, 2).contiguous() v = v.unsqueeze(0).transpose(1, 2).contiguous() - sparse_map, lut, real_topk = get_block_map(q, k, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK) + sparse_map, _, _ = get_block_map(q, k, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK) lut, valid_block_num = block_map_ordinal_lut_triton(sparse_map) out = sage3_block_sparse_attn(q, k, v, lut, valid_block_num, per_block_mean=self.per_block_mean) out = out.transpose(1, 2).reshape(max_seqlen_q, -1) @@ -287,7 +283,7 @@ def apply_fa4( # (L, H, D) -> (B, L, H, D) qt = q.unsqueeze(0).transpose(1, 2).contiguous() kt = k.unsqueeze(0).transpose(1, 2).contiguous() - sparse_map, lut, real_topk = get_block_map(qt, kt, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK) + sparse_map, _, _ = get_block_map(qt, kt, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK) # (L, H, D) -> (B, L, H, D) q = q.unsqueeze(0) @@ -331,7 +327,7 @@ def apply_magi( q_block_map = q_block_map.contiguous() k_block_map = k_block_map.contiguous() - sparse_map, lut, real_topk = get_block_map(q_block_map, k_block_map, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK) + sparse_map, _, _ = get_block_map(q_block_map, k_block_map, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK) seqlen, head_num, head_dim = q.shape q_ranges, k_ranges = self.generate_qk_ranges(sparse_map[0], self.BLKQ, self.BLKK, seqlen) diff --git a/lightx2v/common/ops/attn/utils/sla_util.py b/lightx2v/common/ops/attn/utils/sla_util.py index 4472abde2..80fa3493c 100755 --- a/lightx2v/common/ops/attn/utils/sla_util.py +++ b/lightx2v/common/ops/attn/utils/sla_util.py @@ -44,31 +44,48 @@ def mean_pool(x, BLK): return x_mean -def get_block_lut(q, k, topk_ratio, BLKQ=64, BLKK=64, topk=None): - arg_k = k - torch.mean(k, dim=-2, keepdim=True) # smooth-k technique in SageAttention - pooled_qblocks = mean_pool(q, BLKQ) - pooled_kblocks = mean_pool(arg_k, BLKK) +def _force_local_score_blocks(pooled_score, count): + count = min(count, pooled_score.shape[-1]) + if count <= 0: + return pooled_score + + q_blocks, k_blocks = pooled_score.shape[-2:] + centers = (torch.arange(q_blocks, device=pooled_score.device) * 2 + 1) * k_blocks // (2 * q_blocks) + offsets = torch.arange(count, device=pooled_score.device) - count // 2 + local_indices = (centers[:, None] + offsets[None, :]).clamp(0, k_blocks - 1) + local_indices = local_indices.view(*((1,) * (pooled_score.ndim - 2)), q_blocks, count).expand(*pooled_score.shape[:-2], q_blocks, count) + return pooled_score.scatter(-1, local_indices, torch.finfo(pooled_score.dtype).max) + - # GQA - num_q_heads = q.size(1) - num_kv_heads = k.size(1) - if num_q_heads != num_kv_heads: - assert num_q_heads % num_kv_heads == 0, f"Number of Q heads ({num_q_heads}) must be divisible by number of KV heads ({num_kv_heads})" - repeat_factor = num_q_heads // num_kv_heads - pooled_kblocks = pooled_kblocks.repeat_interleave(repeat_factor, dim=1) +def _expand_kv_heads(pooled_kblocks, num_q_heads): + num_kv_heads = pooled_kblocks.shape[1] + if num_q_heads == num_kv_heads: + return pooled_kblocks - if topk is not None and int(topk) > 16: + assert num_q_heads % num_kv_heads == 0, f"Q heads ({num_q_heads}) must be divisible by KV heads ({num_kv_heads})" + return pooled_kblocks.repeat_interleave(num_q_heads // num_kv_heads, dim=1) + + +def _select_block_lut(pooled_qblocks, pooled_kblocks, topk_ratio, topk, force_local_blocks=0): + pooled_kblocks = _expand_kv_heads(pooled_kblocks, pooled_qblocks.shape[1]) + if topk is not None and topk > 16: pooled_score = pooled_qblocks.float() @ pooled_kblocks.float().transpose(-1, -2) else: pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2) + pooled_score = _force_local_score_blocks(pooled_score, force_local_blocks) k_blocks = pooled_score.shape[-1] - # Match the training router: short sequences still retain one key block. - if topk is None: - topk = int(topk_ratio * k_blocks) - topk = max(1, min(k_blocks, int(topk))) - lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices - return lut, topk + topk = int(topk_ratio * k_blocks) if topk is None else topk + # Short sequences still keep one key block to match the training router. + topk = max(1, min(k_blocks, topk)) + return torch.topk(pooled_score, topk, dim=-1, sorted=False).indices, topk + + +def get_block_lut(q, k, topk_ratio, BLKQ=64, BLKK=64, topk=None): + arg_k = k - torch.mean(k, dim=-2, keepdim=True) # smooth-k technique in SageAttention + pooled_qblocks = mean_pool(q, BLKQ) + pooled_kblocks = mean_pool(arg_k, BLKK) + return _select_block_lut(pooled_qblocks, pooled_kblocks, topk_ratio, topk) def get_block_map(q, k, topk_ratio, BLKQ=64, BLKK=64, topk=None): @@ -181,23 +198,6 @@ def mean_pool_centered_nhd(x, x_mean, BLK): return block_mean -def _force_local_score_blocks(pooled_score, count): - count = min(int(count), int(pooled_score.shape[-1])) - if count <= 0: - return pooled_score - q_blocks = int(pooled_score.shape[-2]) - k_blocks = int(pooled_score.shape[-1]) - centers = (torch.arange(q_blocks, device=pooled_score.device) * 2 + 1) * k_blocks // (2 * q_blocks) - offsets = torch.arange(count, device=pooled_score.device) - count // 2 - local_indices = (centers[:, None] + offsets[None, :]).clamp(0, k_blocks - 1) - local_indices = local_indices.view(*((1,) * (pooled_score.ndim - 2)), q_blocks, count).expand(*pooled_score.shape[:-2], q_blocks, count) - return pooled_score.scatter( - -1, - local_indices, - torch.finfo(pooled_score.dtype).max, - ) - - def _get_block_lut_nhd( q, k, @@ -210,26 +210,7 @@ def _get_block_lut_nhd( pooled_qblocks = mean_pool_nhd(q, BLKQ) k_mean = torch.mean(k, dim=0) pooled_kblocks = mean_pool_centered_nhd(k, k_mean, BLKK) - - num_q_heads = q.size(1) - num_kv_heads = k.size(1) - if num_q_heads != num_kv_heads: - assert num_q_heads % num_kv_heads == 0, f"Number of Q heads ({num_q_heads}) must be divisible by number of KV heads ({num_kv_heads})" - pooled_kblocks = pooled_kblocks.repeat_interleave( - num_q_heads // num_kv_heads, - dim=1, - ) - - if topk is not None and int(topk) > 16: - pooled_score = pooled_qblocks.float() @ pooled_kblocks.float().transpose(-1, -2) - else: - pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2) - pooled_score = _force_local_score_blocks(pooled_score, force_local_blocks) - k_blocks = pooled_score.shape[-1] - if topk is None: - topk = int(topk_ratio * k_blocks) - topk = max(1, min(k_blocks, int(topk))) - return torch.topk(pooled_score, topk, dim=-1, sorted=False).indices, topk + return _select_block_lut(pooled_qblocks, pooled_kblocks, topk_ratio, topk, force_local_blocks) def get_block_lut_nhd_uncentered( @@ -243,26 +224,7 @@ def get_block_lut_nhd_uncentered( ): pooled_qblocks = mean_pool_nhd(q, BLKQ) pooled_kblocks = mean_pool_nhd(k, BLKK) - - num_q_heads = q.size(1) - num_kv_heads = k.size(1) - if num_q_heads != num_kv_heads: - assert num_q_heads % num_kv_heads == 0, f"Number of Q heads ({num_q_heads}) must be divisible by number of KV heads ({num_kv_heads})" - pooled_kblocks = pooled_kblocks.repeat_interleave( - num_q_heads // num_kv_heads, - dim=1, - ) - - if topk is not None and int(topk) > 16: - pooled_score = pooled_qblocks.float() @ pooled_kblocks.float().transpose(-1, -2) - else: - pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2) - pooled_score = _force_local_score_blocks(pooled_score, force_local_blocks) - k_blocks = pooled_score.shape[-1] - if topk is None: - topk = int(topk_ratio * k_blocks) - topk = max(1, min(k_blocks, int(topk))) - return torch.topk(pooled_score, topk, dim=-1, sorted=False).indices, topk + return _select_block_lut(pooled_qblocks, pooled_kblocks, topk_ratio, topk, force_local_blocks) @torch.library.custom_op( @@ -307,12 +269,8 @@ def get_block_lut_nhd( topk=None, force_local_blocks=0, ): - if topk is not None and int(topk) > 16 and int(force_local_blocks) == 0: - topk = int(topk) - return ( - block_lut_nhd(q, k, float(topk_ratio), BLKQ, BLKK, topk), - topk, - ) + if topk is not None and topk > 16 and not force_local_blocks: + return block_lut_nhd(q, k, float(topk_ratio), BLKQ, BLKK, topk), topk return _get_block_lut_nhd( q, k, diff --git a/lightx2v_platform/ops/attn/mthreads_musa/moorcat_sparse.py b/lightx2v_platform/ops/attn/mthreads_musa/moorcat_sparse.py index 3eb6dbc9d..79a0ef7c3 100644 --- a/lightx2v_platform/ops/attn/mthreads_musa/moorcat_sparse.py +++ b/lightx2v_platform/ops/attn/mthreads_musa/moorcat_sparse.py @@ -1,7 +1,3 @@ -from __future__ import annotations - -import math - import torch from loguru import logger @@ -19,24 +15,13 @@ _VALIDATED_CONTRACTS = set() -def _runtime_tensors( - q: torch.Tensor, - seq_len: int, - heads: int, - q_blocks: int, - topk: int, -): +def _runtime_tensors(q: torch.Tensor, seq_len: int, heads: int, q_blocks: int, topk: int): key = (q.device, seq_len, heads, q_blocks, topk) cached = _RUNTIME_CACHE.get(key) if cached is None: cached = ( torch.arange(1, heads + 1, dtype=torch.int32, device=q.device), - torch.full( - (1, heads, q_blocks), - topk, - dtype=torch.int32, - device=q.device, - ), + torch.full((1, heads, q_blocks), topk, dtype=torch.int32, device=q.device), torch.tensor([0, seq_len], dtype=torch.int32, device=q.device), ) _RUNTIME_CACHE[key] = cached @@ -52,18 +37,18 @@ def _validate_lut_once( topk: int, ) -> None: expected = (1, heads, q_blocks, topk) - if tuple(block_indices.shape) != expected: - raise ValueError(f"Moorcat Q128 LUT must have shape {expected}, got {tuple(block_indices.shape)}") + if block_indices.shape != expected: + raise ValueError(f"expected LUT shape {expected}, got {tuple(block_indices.shape)}") contract = (expected, kv_blocks, block_indices.device) if contract in _VALIDATED_CONTRACTS: return min_index = int(block_indices.min().item()) max_index = int(block_indices.max().item()) if min_index < 0 or max_index >= kv_blocks: - raise ValueError(f"Moorcat Q128 LUT indices must be in [0, {kv_blocks}), got [{min_index}, {max_index}]") + raise ValueError(f"LUT index range [{min_index}, {max_index}] exceeds [0, {kv_blocks})") ordered = torch.sort(block_indices, dim=-1).values if not bool(torch.all(ordered[..., 1:] != ordered[..., :-1]).item()): - raise ValueError("Moorcat Q128 LUT rows must contain distinct K blocks") + raise ValueError("LUT rows must not contain duplicate K blocks") _VALIDATED_CONTRACTS.add(contract) @@ -81,12 +66,9 @@ def musa_moorcat_sparse( softmax_scale: float, validate_lut: bool, ) -> torch.Tensor: - if _blocksparse is None: - raise RuntimeError("Moorcat is not importable. Install the Moorcat package supplied with the official MUSA SGL image.") - seq_len, heads, _ = q.shape - q_blocks = math.ceil(seq_len / _BLOCK_SIZE) - kv_blocks = math.ceil(k.shape[0] / _BLOCK_SIZE) + q_blocks = (seq_len + _BLOCK_SIZE - 1) // _BLOCK_SIZE + kv_blocks = (k.shape[0] + _BLOCK_SIZE - 1) // _BLOCK_SIZE if validate_lut: _validate_lut_once( block_indices, @@ -144,13 +126,12 @@ class MusaMoorcatSparseOperator: k_block_size = _BLOCK_SIZE block_indices_only = True - def __init__(self, operator_setting=None): - setting = dict(operator_setting or {}) - self.topk = int(setting.get("topk", 16)) - self.center_k = bool(setting.get("center_k", True)) - self.validate_lut = bool(setting.get("validate_lut", setting.get("validate_mask", True))) - if self.topk <= 0: - raise ValueError(f"Moorcat topk must be positive, got {self.topk}") + def __init__(self, operator_setting): + if _blocksparse is None: + raise RuntimeError("Moorcat blocksparse extension is unavailable") + self.topk = operator_setting.get("topk", 16) + self.center_k = operator_setting.get("center_k", True) + self.validate_lut = operator_setting.get("validate_lut", operator_setting.get("validate_mask", True)) @torch.compiler.disable def __call__( @@ -166,25 +147,6 @@ def __call__( max_seqlen_kv=None, **kwargs, ): - if _blocksparse is None: - raise RuntimeError("The official Moorcat binary is unavailable. Install the Moorcat package supplied with the official MUSA SGL image.") - if q.ndim != 3 or k.ndim != 3 or v.ndim != 3: - raise ValueError("Moorcat requires flattened [tokens, heads, head_dim] Q/K/V") - if q.shape != k.shape or q.shape != v.shape: - raise ValueError(f"Moorcat supports self-attention only, got q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}") - if q.dtype != torch.bfloat16 or k.dtype != q.dtype or v.dtype != q.dtype: - raise ValueError("Moorcat requires BF16 Q/K/V") - if q.shape[-1] != 128: - raise ValueError(f"Moorcat requires head_dim=128, got {q.shape[-1]}") - if bool(kwargs.get("causal", False)): - raise ValueError("Moorcat supports non-causal attention only") - if cu_seqlens_q is not None and cu_seqlens_q.numel() != 2: - raise ValueError("Moorcat supports one packed sequence only") - if cu_seqlens_kv is not None and cu_seqlens_kv.numel() != 2: - raise ValueError("Moorcat supports one packed sequence only") - if block_indices is None: - raise ValueError("Moorcat requires Q128/K128 block indices") - softmax_scale = kwargs.get("softmax_scale") if softmax_scale is None: softmax_scale = q.shape[-1] ** -0.5 From db05d29ea9a9d3c3965dd26ac0fbaec4d7799142 Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Sun, 6 Sep 2026 09:28:54 +0000 Subject: [PATCH 3/5] perf(musa): use SageAttention2 for H3 TP8 inference step latency 0.912378s -> 0.873070s, 1.045x speedup --- configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json index 06ab5f4ed..374ef6067 100644 --- a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json @@ -13,7 +13,8 @@ "vae_decode_parallel": true, "lazy_load": false, "unload_modules": false, - "attn_type": "flash_attn3", + "attn_type": "sage_attn2", + "refiner_attn_type": "flash_attn3", "rms_type": "sgl-kernel", "rope_type": "torch_real_rope", "feature_caching": "NoCaching", From fb4a619f5efa99c457a2f7bf5277968ce2064c88 Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Sun, 6 Sep 2026 09:38:24 +0000 Subject: [PATCH 4/5] feat(musa): add H3 TP8 sparse inference config --- .../minimax_h3_t2av_tp8_sparse.json | 46 +++++++++++++++++++ .../run_minimax_h3_t2av_tp8_sparse.sh | 27 +++++++++++ 2 files changed, 73 insertions(+) create mode 100644 configs/platforms/mthreads_musa/minimax_h3_t2av_tp8_sparse.json create mode 100755 scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp8_sparse.sh diff --git a/configs/platforms/mthreads_musa/minimax_h3_t2av_tp8_sparse.json b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp8_sparse.json new file mode 100644 index 000000000..0aa251cb5 --- /dev/null +++ b/configs/platforms/mthreads_musa/minimax_h3_t2av_tp8_sparse.json @@ -0,0 +1,46 @@ +{ + "infer_steps": 30, + "target_video_length": 124, + "target_height": 544, + "target_width": 960, + "fps": 24, + "target_fps": 24, + "enable_cfg": false, + "cpu_offload": false, + "text_encoder_cpu_offload": false, + "text_encoder_tensor_parallel": true, + "vae_cpu_offload": false, + "vae_decode_parallel": true, + "lazy_load": false, + "unload_modules": false, + "attn_type": "dynamic_sparse_attn", + "refiner_attn_type": "flash_attn3", + "dynamic_sparse_attn_setting": { + "topk": 15, + "force_local_blocks": 2, + "operator": "musa_moorcat_sparse", + "operator_setting": { + "validate_lut": true + } + }, + "rms_type": "sgl-kernel", + "rope_type": "torch_real_rope", + "feature_caching": "NoCaching", + "use_compile": true, + "warmup": true, + "vae_use_compile": true, + "video_flow_shift": 12.0, + "audio_flow_shift": 3.0, + "vae_spatial_scale_factor": 16, + "audio_sampling_rate": 32000, + "audio_latents_per_second": 40, + "audio_channels": 2, + "keep_latents_dtype_in_scheduler": true, + "parallel": { + "tensor_p_size": 8 + }, + "video_codec_options": { + "preset": "ultrafast", + "crf": "18" + } +} diff --git a/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp8_sparse.sh b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp8_sparse.sh new file mode 100755 index 000000000..3056723fb --- /dev/null +++ b/scripts/platforms/mthreads_musa/run_minimax_h3_t2av_tp8_sparse.sh @@ -0,0 +1,27 @@ +#!/bin/bash + +lightx2v_path=/data/wushuo/LightX2V +model_path=/data/MiniMax-H3 + +export PLATFORM=musa +export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 + +source "${lightx2v_path}/scripts/base/base.sh" +export DTYPE=BF16 +export SENSITIVE_LAYER_DTYPE=BF16 + +config_path=${lightx2v_path}/configs/platforms/mthreads_musa/minimax_h3_t2av_tp8_sparse.json +output_path=${lightx2v_path}/save_results/output_lightx2v_minimax_h3_t2av_tp8_sparse.mp4 +log_path=${lightx2v_path}/save_results/minimax_h3_t2av_544p_124_8gpu_tp8_sparse.log + +prompt='In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' + +nohup torchrun --standalone --nproc_per_node=8 -m lightx2v.infer \ + --model_cls minimax_h3 \ + --task t2av \ + --model_path "${model_path}" \ + --config_json "${config_path}" \ + --prompt "${prompt}" \ + --save_result_path "${output_path}" \ + --seed 0 \ + > "${log_path}" 2>&1 & From c69767315435599042261a4a1ecabe1df45573b2 Mon Sep 17 00:00:00 2001 From: WateBear <540295877@qq.com> Date: Tue, 8 Sep 2026 10:09:37 +0000 Subject: [PATCH 5/5] perf(musa): fuse MiniMax-H3 inference ops Fuse Q/K RoPE, AdaLN modulation and gates, SwiGLU, and the legacy RMSNorm scaling path while preserving LightX2V numerical behavior. Handle the non-contiguous AdaLN chunk row stride explicitly. --- lightx2v/common/ops/norm/rms_norm_weight.py | 11 + lightx2v/common/ops/norm/triton_ops.py | 45 +++ .../networks/minimax_h3/infer/pre_infer.py | 2 + .../minimax_h3/infer/transformer_infer.py | 39 ++- .../networks/minimax_h3/infer/triton_ops.py | 291 +++++++++++++++++- .../attn/mthreads_musa/csrc/h3_qk_rope.cpp | 18 ++ .../ops/attn/mthreads_musa/csrc/h3_qk_rope.mu | 212 +++++++++++++ .../ops/attn/mthreads_musa/h3_fused_ops.py | 27 ++ 8 files changed, 619 insertions(+), 26 deletions(-) create mode 100644 lightx2v_platform/ops/attn/mthreads_musa/csrc/h3_qk_rope.cpp create mode 100644 lightx2v_platform/ops/attn/mthreads_musa/csrc/h3_qk_rope.mu create mode 100644 lightx2v_platform/ops/attn/mthreads_musa/h3_fused_ops.py diff --git a/lightx2v/common/ops/norm/rms_norm_weight.py b/lightx2v/common/ops/norm/rms_norm_weight.py index 67503ad1a..a55b0d2c2 100755 --- a/lightx2v/common/ops/norm/rms_norm_weight.py +++ b/lightx2v/common/ops/norm/rms_norm_weight.py @@ -10,6 +10,7 @@ fused_qk_norm_3drope, fused_qk_rms_norm, rms_norm_kernel, + rms_norm_legacy_scale, ) from lightx2v.common.ops.utils import * from lightx2v.utils.envs import * @@ -375,6 +376,16 @@ def apply(self, input_tensor): return input_tensor +@RMS_WEIGHT_REGISTER("h3_legacy_triton") +class RMSWeightH3LegacyTriton(RMSWeight): + def apply(self, input_tensor): + weight = self._get_actual_weight() + if weight is None or self.sensitive_layer_dtype != self.infer_dtype: + return super().apply(input_tensor) + inverse_rms = torch.rsqrt(input_tensor.pow(2).mean(-1, keepdim=True) + self.eps) + return rms_norm_legacy_scale(input_tensor, weight, inverse_rms) + + @RMS_WEIGHT_REGISTER("fp32_variance") class RMSWeightFP32(RMSWeight): def __init__( diff --git a/lightx2v/common/ops/norm/triton_ops.py b/lightx2v/common/ops/norm/triton_ops.py index 5979f6c94..940c90b69 100644 --- a/lightx2v/common/ops/norm/triton_ops.py +++ b/lightx2v/common/ops/norm/triton_ops.py @@ -992,6 +992,51 @@ def rms_norm_kernel( return y +@triton.jit +def _rms_norm_legacy_scale_kernel( + output, + x, + inverse_rms, + weight, + n_elements, + HIDDEN_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + columns = offsets % HIDDEN_SIZE + rows = offsets // HIDDEN_SIZE + values = tl.load(x + offsets, mask=mask) + scales = tl.load(inverse_rms + rows, mask=mask) + weights = tl.load(weight + columns, mask=mask) + normalized = (values * scales).to(tl.bfloat16) + tl.store(output + offsets, (normalized * weights).to(tl.bfloat16), mask=mask) + + +def rms_norm_legacy_scale( + x: torch.Tensor, + weight: torch.Tensor, + inverse_rms: torch.Tensor, +) -> torch.Tensor: + x = x.contiguous() + output = torch.empty_like(x) + block_size = 1024 + grid = (triton.cdiv(x.numel(), block_size),) + device_module = torch.cuda if x.device.type == "cuda" else torch.musa + with device_module.device(x.device): + torch.library.wrap_triton(_rms_norm_legacy_scale_kernel)[grid]( + output, + x, + inverse_rms, + weight, + x.numel(), + HIDDEN_SIZE=x.shape[-1], + BLOCK_SIZE=block_size, + num_warps=8, + ) + return output + + @triton.jit def _fused_qk_rms_norm_kernel( q_ptr, diff --git a/lightx2v/models/networks/minimax_h3/infer/pre_infer.py b/lightx2v/models/networks/minimax_h3/infer/pre_infer.py index 2ccd86806..95a225a83 100644 --- a/lightx2v/models/networks/minimax_h3/infer/pre_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/pre_infer.py @@ -91,6 +91,8 @@ def _rotary_embedding(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, t freqs_t, freqs_h, freqs_w = freqs.unbind(dim=1) freqs = torch.cat((freqs_t, freqs_h, freqs_w), dim=-1) freqs = torch.cat((freqs, freqs), dim=-1) + if self.config.get("rope_type") == "minimax_h3_musa_rope_bf16": + return freqs.cos().to(GET_DTYPE()), freqs.sin().to(GET_DTYPE()) return freqs.cos(), freqs.sin() def infer(self, weights, prompt_embeds): diff --git a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py index f7366425f..c09a7eecf 100644 --- a/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py +++ b/lightx2v/models/networks/minimax_h3/infer/transformer_infer.py @@ -3,6 +3,11 @@ import torch.nn.functional as F from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer +from lightx2v.models.networks.minimax_h3.infer.triton_ops import ( + apply_h3_gate_triton, + apply_h3_scale_shift_triton, + apply_h3_swiglu_triton, +) from lightx2v.utils.envs import GET_DTYPE @@ -36,6 +41,8 @@ def __init__(self, config): self.seq_p_group = None self.infer_func = self.infer_without_offload self.use_adaln_cache = bool(config.get("use_adaln_cache", False)) + self.use_triton_modulation = bool(config.get("use_triton_modulation", False)) + self.use_triton_swiglu = bool(config.get("use_triton_swiglu", False)) self._adaln_cache = {} self._current_adaln_tables = None self._adaln_cache_hit = False @@ -102,10 +109,24 @@ def _attention(self, weights, hidden_states, pre_infer_out): out = torch.cat((aux_out, out), dim=0) return weights.to_out.apply(out.to(self.infer_dtype)) - @staticmethod - def _ff(weights, hidden_states): - value, gate = weights.in_proj.apply(hidden_states).chunk(2, dim=-1) - return weights.out_proj.apply(value * F.silu(gate)) + def _ff(self, weights, hidden_states): + projected = weights.in_proj.apply(hidden_states) + if self.use_triton_swiglu: + activated = apply_h3_swiglu_triton(projected) + else: + value, gate = projected.chunk(2, dim=-1) + activated = value * F.silu(gate) + return weights.out_proj.apply(activated) + + def _scale_shift(self, hidden_states, shift, scale, indices): + if self.use_triton_modulation: + return apply_h3_scale_shift_triton(hidden_states, shift, scale, indices) + return hidden_states * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices) + + def _gate(self, residual, gate, hidden_states, indices): + if self.use_triton_modulation: + return apply_h3_gate_triton(residual, gate, hidden_states, indices) + return residual + gate.index_select(0, indices) * hidden_states def infer_block(self, weights, hidden_states, pre_infer_out, modulation=None): # Keep the Python cache lookup outside the compiled block. @@ -116,15 +137,13 @@ def infer_block(self, weights, hidden_states, pre_infer_out, modulation=None): residual = hidden_states normed = weights.norm1.apply(hidden_states) - normed = normed * (1.0 + scale_msa.index_select(0, indices)) - normed = normed + shift_msa.index_select(0, indices) - hidden_states = residual + gate_msa.index_select(0, indices) * self._attention(weights.attn, normed, pre_infer_out) + normed = self._scale_shift(normed, shift_msa, scale_msa, indices) + hidden_states = self._gate(residual, gate_msa, self._attention(weights.attn, normed, pre_infer_out), indices) residual = hidden_states normed = weights.norm2.apply(hidden_states) - normed = normed * (1.0 + scale_mlp.index_select(0, indices)) - normed = normed + shift_mlp.index_select(0, indices) - hidden_states = residual + gate_mlp.index_select(0, indices) * self._ff(weights.ff, normed) + normed = self._scale_shift(normed, shift_mlp, scale_mlp, indices) + hidden_states = self._gate(residual, gate_mlp, self._ff(weights.ff, normed), indices) return hidden_states def _compute_adaln_table(self, weights, pre_infer_out): diff --git a/lightx2v/models/networks/minimax_h3/infer/triton_ops.py b/lightx2v/models/networks/minimax_h3/infer/triton_ops.py index 1d4c618d0..5cbe47f5a 100644 --- a/lightx2v/models/networks/minimax_h3/infer/triton_ops.py +++ b/lightx2v/models/networks/minimax_h3/infer/triton_ops.py @@ -16,9 +16,9 @@ class MiniMaxH3TritonRope(RopeTemplate): """Partial split-half RoPE used by MiniMax-H3. H3 rotates only the leading RoPE dimensions of each attention head and - leaves the remaining channels unchanged. CUDA tensors use the local - Triton kernel; other devices (or environments without Triton) fall back - to the shared real-valued RoPE implementation with identical layout. + leaves the remaining channels unchanged. CUDA and MUSA tensors use the + local Triton kernel; other devices (or environments without Triton) fall + back to the shared real-valued RoPE implementation with identical layout. """ def __init__(self, layout="split_half", compute_dtype=torch.float32): @@ -30,21 +30,42 @@ def __init__(self, layout="split_half", compute_dtype=torch.float32): def apply(self, query, key, freqs, rotary_dim=None, **kwargs): cos, sin = freqs rotary_dim = cos.shape[-1] if rotary_dim is None else rotary_dim - if query.is_cuda and key.is_cuda and triton is not None: - return ( - apply_partial_split_half_rotary_triton(query, cos, sin, rotary_dim), - apply_partial_split_half_rotary_triton(key, cos, sin, rotary_dim), - ) + if query.device.type in {"cuda", "musa"} and key.device == query.device and triton is not None: + return apply_partial_split_half_qk_rotary_triton(query, key, cos, sin, rotary_dim) return self.torch_rope.apply(query, key, freqs, rotary_dim=rotary_dim, **kwargs) def apply_single(self, x, freqs, rotary_dim=None, **kwargs): cos, sin = freqs rotary_dim = cos.shape[-1] if rotary_dim is None else rotary_dim - if x.is_cuda and triton is not None: + if x.device.type in {"cuda", "musa"} and triton is not None: return apply_partial_split_half_rotary_triton(x, cos, sin, rotary_dim) return self.torch_rope.apply_single(x, freqs, rotary_dim=rotary_dim, **kwargs) +@ROPE_REGISTER("minimax_h3_musa_rope") +class MiniMaxH3MusaRope(MiniMaxH3TritonRope): + def apply(self, query, key, freqs, rotary_dim=None, **kwargs): + cos, sin = freqs + rotary_dim = cos.shape[-1] if rotary_dim is None else rotary_dim + if query.device.type == "musa" and rotary_dim == 96: + from lightx2v_platform.ops.attn.mthreads_musa.h3_fused_ops import apply_h3_qk_rope_fp32 + + return apply_h3_qk_rope_fp32(query, key, cos, sin) + return super().apply(query, key, freqs, rotary_dim=rotary_dim, **kwargs) + + +@ROPE_REGISTER("minimax_h3_musa_rope_bf16") +class MiniMaxH3MusaBf16Rope(MiniMaxH3TritonRope): + def apply(self, query, key, freqs, rotary_dim=None, **kwargs): + cos, sin = freqs + rotary_dim = cos.shape[-1] if rotary_dim is None else rotary_dim + if query.device.type == "musa" and rotary_dim == 96: + from lightx2v_platform.ops.attn.mthreads_musa.h3_fused_ops import apply_h3_qk_rope + + return apply_h3_qk_rope(query, key, cos, sin) + return super().apply(query, key, freqs, rotary_dim=rotary_dim, **kwargs) + + if triton is not None: @triton.jit @@ -62,8 +83,9 @@ def _partial_split_half_rotary_kernel( ROTARY_DIM: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): - row_idx = tl.program_id(0) - token_idx = (row_idx // num_heads) % num_tokens + token_idx = tl.program_id(0) + head_idx = tl.program_id(1) + row_idx = token_idx * num_heads + head_idx offsets = tl.arange(0, BLOCK_SIZE) head_mask = offsets < HEAD_SIZE @@ -88,6 +110,138 @@ def _partial_split_half_rotary_kernel( output = tl.where(rotary_mask, rotated_output, x_fp32) tl.store(output_row_ptr + offsets, output.to(x.dtype), mask=head_mask) + @triton.jit + def _partial_split_half_qk_rotary_kernel( + query_output_ptr, + key_output_ptr, + query_ptr, + key_ptr, + cos_ptr, + sin_ptr, + num_heads, + stride_x_row, + stride_cos_row, + stride_sin_row, + HEAD_SIZE: tl.constexpr, + ROTARY_DIM: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + ): + token_idx = tl.program_id(0) + head_idx = tl.program_id(1) + row_idx = token_idx * num_heads + head_idx + + offsets = tl.arange(0, BLOCK_SIZE) + head_mask = offsets < HEAD_SIZE + rotary_mask = offsets < ROTARY_DIM + rotary_half = ROTARY_DIM // 2 + partner_offsets = tl.where(offsets < rotary_half, offsets + rotary_half, offsets - rotary_half) + + query_row_ptr = query_ptr + row_idx * stride_x_row + key_row_ptr = key_ptr + row_idx * stride_x_row + query_output_row_ptr = query_output_ptr + row_idx * stride_x_row + key_output_row_ptr = key_output_ptr + row_idx * stride_x_row + cos_row_ptr = cos_ptr + token_idx * stride_cos_row + sin_row_ptr = sin_ptr + token_idx * stride_sin_row + + cos = tl.load(cos_row_ptr + offsets, mask=rotary_mask, other=1.0) + sin = tl.load(sin_row_ptr + offsets, mask=rotary_mask, other=0.0) + + query = tl.load(query_row_ptr + offsets, mask=head_mask, other=0.0) + query_partner = tl.load(query_row_ptr + partner_offsets, mask=rotary_mask, other=0.0) + query_cos = (query.to(tl.float32) * cos.to(tl.float32)).to(tl.bfloat16) + query_partner_sin = (query_partner.to(tl.float32) * sin.to(tl.float32)).to(tl.bfloat16) + query_partner_sin = tl.where(offsets < rotary_half, -query_partner_sin, query_partner_sin) + query_rotated = (query_cos.to(tl.float32) + query_partner_sin.to(tl.float32)).to(tl.bfloat16) + query_result = tl.where(rotary_mask, query_rotated, query) + tl.store(query_output_row_ptr + offsets, query_result, mask=head_mask) + + key = tl.load(key_row_ptr + offsets, mask=head_mask, other=0.0) + key_partner = tl.load(key_row_ptr + partner_offsets, mask=rotary_mask, other=0.0) + key_cos = (key.to(tl.float32) * cos.to(tl.float32)).to(tl.bfloat16) + key_partner_sin = (key_partner.to(tl.float32) * sin.to(tl.float32)).to(tl.bfloat16) + key_partner_sin = tl.where(offsets < rotary_half, -key_partner_sin, key_partner_sin) + key_rotated = (key_cos.to(tl.float32) + key_partner_sin.to(tl.float32)).to(tl.bfloat16) + key_result = tl.where(rotary_mask, key_rotated, key) + tl.store(key_output_row_ptr + offsets, key_result, mask=head_mask) + + @triton.jit + def _h3_scale_shift_kernel( + output, + x, + shift, + scale, + indices, + shift_row_stride, + scale_row_stride, + n_elements, + HIDDEN_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + ): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + rows = offsets // HIDDEN_SIZE + columns = offsets - rows * HIDDEN_SIZE + parameter_rows = tl.load(indices + rows, mask=mask, other=0) + + values = tl.load(x + offsets, mask=mask) + shifts = tl.load( + shift + parameter_rows * shift_row_stride + columns, + mask=mask, + ) + scales = tl.load( + scale + parameter_rows * scale_row_stride + columns, + mask=mask, + ) + factors = (1.0 + scales).to(tl.bfloat16) + scaled = (values * factors).to(tl.bfloat16) + tl.store(output + offsets, (scaled + shifts).to(tl.bfloat16), mask=mask) + + @triton.jit + def _h3_gate_kernel( + output, + residual, + gate, + branch, + indices, + gate_row_stride, + n_elements, + HIDDEN_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + ): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + rows = offsets // HIDDEN_SIZE + columns = offsets - rows * HIDDEN_SIZE + parameter_rows = tl.load(indices + rows, mask=mask, other=0) + + residual_values = tl.load(residual + offsets, mask=mask) + gate_values = tl.load( + gate + parameter_rows * gate_row_stride + columns, + mask=mask, + ) + branch_values = tl.load(branch + offsets, mask=mask) + product = (gate_values * branch_values).to(tl.bfloat16) + tl.store(output + offsets, (residual_values + product).to(tl.bfloat16), mask=mask) + + @triton.jit + def _h3_swiglu_kernel( + output, + projected, + HIDDEN_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + ): + row = tl.program_id(0) + columns = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = columns < HIDDEN_SIZE + row_offset = row * 2 * HIDDEN_SIZE + + value = tl.load(projected + row_offset + columns, mask=mask) + gate = tl.load(projected + row_offset + HIDDEN_SIZE + columns, mask=mask) + sigmoid = (1.0 / (1.0 + tl.exp2(-gate.to(tl.float32) * 1.4426950408889634))).to(tl.bfloat16) + activated = (gate.to(tl.float32) * sigmoid.to(tl.float32)).to(tl.bfloat16) + result = (activated.to(tl.float32) * value.to(tl.float32)).to(tl.bfloat16) + tl.store(output + row * HIDDEN_SIZE + columns, result, mask=mask) + def apply_partial_split_half_rotary_triton( x: torch.Tensor, @@ -97,8 +251,8 @@ def apply_partial_split_half_rotary_triton( ) -> torch.Tensor: if triton is None: raise RuntimeError("Triton is required for MiniMax-H3 Triton RoPE") - if not x.is_cuda: - raise ValueError("MiniMax-H3 Triton RoPE requires a CUDA tensor") + if x.device.type not in {"cuda", "musa"}: + raise ValueError("MiniMax-H3 Triton RoPE requires a CUDA or MUSA tensor") if x.ndim != 3: raise ValueError(f"MiniMax-H3 Triton RoPE expects [L, H, D], got {tuple(x.shape)}") if cos.shape != sin.shape or cos.ndim != 2: @@ -118,9 +272,10 @@ def apply_partial_split_half_rotary_triton( sin = sin.to(device=x.device).contiguous() output = torch.empty_like(x) block_size = triton.next_power_of_2(head_size) - grid = (num_tokens * num_heads,) - with torch.cuda.device(x.device): - _partial_split_half_rotary_kernel[grid]( + grid = (num_tokens, num_heads) + device_module = torch.cuda if x.device.type == "cuda" else torch.musa + with device_module.device(x.device): + torch.library.wrap_triton(_partial_split_half_rotary_kernel)[grid]( output, x, cos, @@ -135,3 +290,107 @@ def apply_partial_split_half_rotary_triton( BLOCK_SIZE=block_size, ) return output + + +def apply_partial_split_half_qk_rotary_triton( + query: torch.Tensor, + key: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + rotary_dim: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + if query.shape != key.shape or query.device != key.device or query.dtype != key.dtype: + raise ValueError("query and key must have the same shape, device, and dtype") + if query.ndim != 3: + raise ValueError(f"MiniMax-H3 Triton RoPE expects [L, H, D], got {tuple(query.shape)}") + if cos.shape != sin.shape or cos.ndim != 2: + raise ValueError(f"cos and sin must have matching [L, R] shapes, got {tuple(cos.shape)} and {tuple(sin.shape)}") + + num_tokens, num_heads, head_size = query.shape + rotary_dim = cos.shape[-1] if rotary_dim is None else int(rotary_dim) + if cos.shape[0] != num_tokens or rotary_dim != cos.shape[-1]: + raise ValueError("RoPE frequencies do not match the query/key shape") + if rotary_dim <= 0 or rotary_dim > head_size or rotary_dim % 2: + raise ValueError(f"rotary_dim must be positive, even, and <= head_size; got rotary_dim={rotary_dim}, head_size={head_size}") + + query = query.contiguous() + key = key.contiguous() + cos = cos.to(device=query.device).contiguous() + sin = sin.to(device=query.device).contiguous() + query_output = torch.empty_like(query) + key_output = torch.empty_like(key) + grid = (num_tokens, num_heads) + block_size = triton.next_power_of_2(head_size) + device_module = torch.cuda if query.device.type == "cuda" else torch.musa + with device_module.device(query.device): + torch.library.wrap_triton(_partial_split_half_qk_rotary_kernel)[grid]( + query_output, + key_output, + query, + key, + cos, + sin, + num_heads, + query.stride(1), + cos.stride(0), + sin.stride(0), + HEAD_SIZE=head_size, + ROTARY_DIM=rotary_dim, + BLOCK_SIZE=block_size, + ) + return query_output, key_output + + +def apply_h3_scale_shift_triton(x, shift, scale, indices): + output = torch.empty_like(x) + block_size = 1024 + grid = (triton.cdiv(x.numel(), block_size),) + torch.library.wrap_triton(_h3_scale_shift_kernel)[grid]( + output, + x, + shift, + scale, + indices, + shift.stride(0), + scale.stride(0), + x.numel(), + HIDDEN_SIZE=x.shape[-1], + BLOCK_SIZE=block_size, + num_warps=8, + ) + return output + + +def apply_h3_gate_triton(residual, gate, branch, indices): + output = torch.empty_like(residual) + block_size = 1024 + grid = (triton.cdiv(residual.numel(), block_size),) + torch.library.wrap_triton(_h3_gate_kernel)[grid]( + output, + residual, + gate, + branch, + indices, + gate.stride(0), + residual.numel(), + HIDDEN_SIZE=residual.shape[-1], + BLOCK_SIZE=block_size, + num_warps=8, + ) + return output + + +def apply_h3_swiglu_triton(projected): + hidden_size = projected.shape[-1] // 2 + projected = projected.view(-1, 2 * hidden_size) + output = torch.empty((projected.shape[0], hidden_size), dtype=projected.dtype, device=projected.device) + block_size = 1024 + grid = (projected.shape[0], triton.cdiv(hidden_size, block_size)) + torch.library.wrap_triton(_h3_swiglu_kernel)[grid]( + output, + projected, + HIDDEN_SIZE=hidden_size, + BLOCK_SIZE=block_size, + num_warps=8, + ) + return output diff --git a/lightx2v_platform/ops/attn/mthreads_musa/csrc/h3_qk_rope.cpp b/lightx2v_platform/ops/attn/mthreads_musa/csrc/h3_qk_rope.cpp new file mode 100644 index 000000000..c1b98c1c4 --- /dev/null +++ b/lightx2v_platform/ops/attn/mthreads_musa/csrc/h3_qk_rope.cpp @@ -0,0 +1,18 @@ +#include + +std::vector h3_qk_rope( + torch::Tensor query, + torch::Tensor key, + torch::Tensor cos, + torch::Tensor sin); + +std::vector h3_qk_rope_fp32( + torch::Tensor query, + torch::Tensor key, + torch::Tensor cos, + torch::Tensor sin); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("h3_qk_rope", &h3_qk_rope); + module.def("h3_qk_rope_fp32", &h3_qk_rope_fp32); +} diff --git a/lightx2v_platform/ops/attn/mthreads_musa/csrc/h3_qk_rope.mu b/lightx2v_platform/ops/attn/mthreads_musa/csrc/h3_qk_rope.mu new file mode 100644 index 000000000..4070167f2 --- /dev/null +++ b/lightx2v_platform/ops/attn/mthreads_musa/csrc/h3_qk_rope.mu @@ -0,0 +1,212 @@ +#include +#include + +#include "torch_musa/csrc/core/MUSAException.h" +#include "torch_musa/csrc/core/MUSAGuard.h" +#include "torch_musa/csrc/core/MUSAStream.h" + +namespace { + +constexpr int kHeadSize = 128; +constexpr int kRotarySize = 96; +constexpr int kRotaryHalf = 48; + +__device__ __forceinline__ __mt_bfloat16 multiply( + __mt_bfloat16 left, + __mt_bfloat16 right) { + return __float2bfloat16_rn(__bfloat162float(left) * __bfloat162float(right)); +} + +__device__ __forceinline__ __mt_bfloat16 add( + __mt_bfloat16 left, + __mt_bfloat16 right) { + return __float2bfloat16_rn(__bfloat162float(left) + __bfloat162float(right)); +} + +__device__ __forceinline__ __mt_bfloat16 subtract( + __mt_bfloat16 left, + __mt_bfloat16 right) { + return __float2bfloat16_rn(__bfloat162float(left) - __bfloat162float(right)); +} + +__global__ void h3_qk_rope_kernel( + const __mt_bfloat16* query, + const __mt_bfloat16* key, + const __mt_bfloat16* cos, + const __mt_bfloat16* sin, + __mt_bfloat16* query_out, + __mt_bfloat16* key_out, + int heads, + int query_token_stride, + int query_head_stride, + int key_token_stride, + int key_head_stride) { + const int token = blockIdx.x; + const int head = blockIdx.y; + const int column = threadIdx.x; + const int output_offset = (token * heads + head) * kHeadSize + column; + const int query_offset = token * query_token_stride + head * query_head_stride; + const int key_offset = token * key_token_stride + head * key_head_stride; + + if (column >= kRotarySize) { + query_out[output_offset] = query[query_offset + column]; + key_out[output_offset] = key[key_offset + column]; + return; + } + + const int left_column = column < kRotaryHalf ? column : column - kRotaryHalf; + const int right_column = left_column + kRotaryHalf; + const __mt_bfloat16 frequency_cos = cos[token * kRotarySize + column]; + const __mt_bfloat16 frequency_sin = sin[token * kRotarySize + column]; + const __mt_bfloat16 query_left = query[query_offset + left_column]; + const __mt_bfloat16 query_right = query[query_offset + right_column]; + const __mt_bfloat16 key_left = key[key_offset + left_column]; + const __mt_bfloat16 key_right = key[key_offset + right_column]; + + if (column < kRotaryHalf) { + query_out[output_offset] = subtract( + multiply(query_left, frequency_cos), + multiply(query_right, frequency_sin)); + key_out[output_offset] = subtract( + multiply(key_left, frequency_cos), + multiply(key_right, frequency_sin)); + } else { + query_out[output_offset] = add( + multiply(query_left, frequency_sin), + multiply(query_right, frequency_cos)); + key_out[output_offset] = add( + multiply(key_left, frequency_sin), + multiply(key_right, frequency_cos)); + } +} + +__global__ void h3_qk_rope_fp32_kernel( + const __mt_bfloat16* query, + const __mt_bfloat16* key, + const float* cos, + const float* sin, + __mt_bfloat16* query_out, + __mt_bfloat16* key_out, + int heads, + int query_token_stride, + int query_head_stride, + int key_token_stride, + int key_head_stride) { + const int token = blockIdx.x; + const int head = blockIdx.y; + const int column = threadIdx.x; + const int output_offset = (token * heads + head) * kHeadSize + column; + const int query_offset = token * query_token_stride + head * query_head_stride; + const int key_offset = token * key_token_stride + head * key_head_stride; + + if (column >= kRotarySize) { + query_out[output_offset] = query[query_offset + column]; + key_out[output_offset] = key[key_offset + column]; + return; + } + + const int partner_column = column < kRotaryHalf ? column + kRotaryHalf : column - kRotaryHalf; + const float frequency_cos = cos[token * kRotarySize + column]; + const float frequency_sin = sin[token * kRotarySize + column]; + const float query_value = __bfloat162float(query[query_offset + column]); + const float query_partner = __bfloat162float(query[query_offset + partner_column]); + const float key_value = __bfloat162float(key[key_offset + column]); + const float key_partner = __bfloat162float(key[key_offset + partner_column]); + const float sign = column < kRotaryHalf ? -1.0f : 1.0f; + const float query_cos = query_value * frequency_cos; + const float query_sin = sign * query_partner * frequency_sin; + const float key_cos = key_value * frequency_cos; + const float key_sin = sign * key_partner * frequency_sin; + query_out[output_offset] = __float2bfloat16_rn(query_cos + query_sin); + key_out[output_offset] = __float2bfloat16_rn(key_cos + key_sin); +} + +} // namespace + +std::vector h3_qk_rope( + torch::Tensor query, + torch::Tensor key, + torch::Tensor cos, + torch::Tensor sin) { + TORCH_CHECK(query.is_musa() && key.is_musa(), "query and key must be MUSA tensors"); + TORCH_CHECK(query.scalar_type() == torch::kBFloat16, "query must be BF16"); + TORCH_CHECK(key.scalar_type() == torch::kBFloat16, "key must be BF16"); + TORCH_CHECK(cos.scalar_type() == torch::kBFloat16, "cos must be BF16"); + TORCH_CHECK(sin.scalar_type() == torch::kBFloat16, "sin must be BF16"); + TORCH_CHECK(query.dim() == 3 && query.size(2) == kHeadSize, "query must be [T,H,128]"); + TORCH_CHECK(key.sizes() == query.sizes(), "key shape must match query"); + TORCH_CHECK( + cos.dim() == 2 && cos.size(0) == query.size(0) && cos.size(1) == kRotarySize, + "cos must be [T,96]"); + TORCH_CHECK(sin.sizes() == cos.sizes(), "sin shape must match cos"); + TORCH_CHECK(query.stride(2) == 1 && key.stride(2) == 1, "head dimension must be contiguous"); + TORCH_CHECK(cos.is_contiguous() && sin.is_contiguous(), "cos and sin must be contiguous"); + + c10::musa::MUSAGuard device_guard(query.device()); + auto query_out = torch::empty_like(query, torch::MemoryFormat::Contiguous); + auto key_out = torch::empty_like(key, torch::MemoryFormat::Contiguous); + const int rows = query.size(0); + const int heads = query.size(1); + h3_qk_rope_kernel<<< + dim3(rows, heads, 1), + dim3(kHeadSize, 1, 1), + 0, + c10::musa::getCurrentMUSAStream()>>>( + reinterpret_cast(query.data_ptr()), + reinterpret_cast(key.data_ptr()), + reinterpret_cast(cos.data_ptr()), + reinterpret_cast(sin.data_ptr()), + reinterpret_cast<__mt_bfloat16*>(query_out.data_ptr()), + reinterpret_cast<__mt_bfloat16*>(key_out.data_ptr()), + heads, + query.stride(0), + query.stride(1), + key.stride(0), + key.stride(1)); + TORCH_MUSA_CHECK(musaGetLastError()); + return {query_out, key_out}; +} + +std::vector h3_qk_rope_fp32( + torch::Tensor query, + torch::Tensor key, + torch::Tensor cos, + torch::Tensor sin) { + TORCH_CHECK(query.is_musa() && key.is_musa(), "query and key must be MUSA tensors"); + TORCH_CHECK(query.scalar_type() == torch::kBFloat16, "query must be BF16"); + TORCH_CHECK(key.scalar_type() == torch::kBFloat16, "key must be BF16"); + TORCH_CHECK(cos.scalar_type() == torch::kFloat32, "cos must be FP32"); + TORCH_CHECK(sin.scalar_type() == torch::kFloat32, "sin must be FP32"); + TORCH_CHECK(query.dim() == 3 && query.size(2) == kHeadSize, "query must be [T,H,128]"); + TORCH_CHECK(key.sizes() == query.sizes(), "key shape must match query"); + TORCH_CHECK( + cos.dim() == 2 && cos.size(0) == query.size(0) && cos.size(1) == kRotarySize, + "cos must be [T,96]"); + TORCH_CHECK(sin.sizes() == cos.sizes(), "sin shape must match cos"); + TORCH_CHECK(query.stride(2) == 1 && key.stride(2) == 1, "head dimension must be contiguous"); + TORCH_CHECK(cos.is_contiguous() && sin.is_contiguous(), "cos and sin must be contiguous"); + + c10::musa::MUSAGuard device_guard(query.device()); + auto query_out = torch::empty_like(query, torch::MemoryFormat::Contiguous); + auto key_out = torch::empty_like(key, torch::MemoryFormat::Contiguous); + const int rows = query.size(0); + const int heads = query.size(1); + h3_qk_rope_fp32_kernel<<< + dim3(rows, heads, 1), + dim3(kHeadSize, 1, 1), + 0, + c10::musa::getCurrentMUSAStream()>>>( + reinterpret_cast(query.data_ptr()), + reinterpret_cast(key.data_ptr()), + cos.data_ptr(), + sin.data_ptr(), + reinterpret_cast<__mt_bfloat16*>(query_out.data_ptr()), + reinterpret_cast<__mt_bfloat16*>(key_out.data_ptr()), + heads, + query.stride(0), + query.stride(1), + key.stride(0), + key.stride(1)); + TORCH_MUSA_CHECK(musaGetLastError()); + return {query_out, key_out}; +} diff --git a/lightx2v_platform/ops/attn/mthreads_musa/h3_fused_ops.py b/lightx2v_platform/ops/attn/mthreads_musa/h3_fused_ops.py new file mode 100644 index 000000000..13bf54c88 --- /dev/null +++ b/lightx2v_platform/ops/attn/mthreads_musa/h3_fused_ops.py @@ -0,0 +1,27 @@ +from functools import lru_cache +from pathlib import Path + + +@lru_cache(maxsize=1) +def _load_h3_qk_rope(): + from torch_musa.utils.musa_extension import load + + source_dir = Path(__file__).with_name("csrc") + return load( + name="lightx2v_h3_qk_rope", + sources=[ + str(source_dir / "h3_qk_rope.cpp"), + str(source_dir / "h3_qk_rope.mu"), + ], + extra_cflags=["-O3"], + extra_musa_cflags=["-O3", "-ffp-contract=off"], + verbose=False, + ) + + +def apply_h3_qk_rope(query, key, cos, sin): + return _load_h3_qk_rope().h3_qk_rope(query, key, cos, sin) + + +def apply_h3_qk_rope_fp32(query, key, cos, sin): + return _load_h3_qk_rope().h3_qk_rope_fp32(query, key, cos, sin)