sinkhorn 算子

概述

sinkhorn 算子是一个DS系列模型提出的新算法，当前基于 Triton 实现VV算子融合。该算子充分利用了昇腾 NPU 的并行计算能力，通过 Triton 语言编写的内核实现了高性能的计算操作，算子逻辑如下：

def torch_golden(
    mixes: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    hc_mult: int = 4,
    sinkhorn_iters: int = 20,
    eps: float = 1e-6,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    PyTorch native reference implementation of the HC-Split Sinkhorn operator
    Fully aligned with the original operator logic for precision validation of Triton/TileLang versions

    Args:
        mixes: Input tensor, shape [b, s, (2+hc_mult)*hc_mult]
               i.e., [b, s, hc_mult + hc_mult + hc_mult*hc_mult]
        hc_scale: Scale tensor, shape [3]
        hc_base: Bias tensor, shape [(2+hc_mult)*hc_mult]
                 i.e., [hc_mult + hc_mult + hc_mult*hc_mult]
        hc_mult: HC dimension size, default 4
        sinkhorn_iters: Number of Sinkhorn iterations, default 20
        eps: Small constant to prevent division by zero, default 1e-6

    Returns:
        pre: [b, s, hc_mult], Sigmoid activation + eps
        post: [b, s, hc_mult], 2×Sigmoid activation
        comb: [b, s, hc_mult, hc_mult], Sinkhorn-normalized matrix
    """
    # 1. Save original shape for final reshaping
    b, s, _ = mixes.shape
    # Flatten to [b*s, (2+hc_mult)*hc_mult] (align with original operator flattening logic)
    mixes_flat = mixes.view(-1, (2 + hc_mult) * hc_mult)

    # 2. Compute pre: [b*s, hc_mult]
    pre_slice = mixes_flat[:, :hc_mult]  # First hc_mult dimensions
    pre_flat = torch.sigmoid(pre_slice * hc_scale[0] + hc_base[:hc_mult]) + eps

    # 3. Compute post: [b*s, hc_mult]
    post_slice = mixes_flat[:, hc_mult : 2 * hc_mult]  # Middle hc_mult dimensions
    post_flat = 2 * torch.sigmoid(post_slice * hc_scale[1] + hc_base[hc_mult : 2 * hc_mult])

    # 4. Compute initial comb values: [b*s, hc_mult, hc_mult]
    comb_slice = mixes_flat[:, 2 * hc_mult :]  # Last hc_mult×hc_mult dimensions
    comb_init_flat = comb_slice.view(-1, hc_mult, hc_mult)  # Reshape to 2D matrix
    # Linear transformation (scale + base): base is broadcasted to batch dimension
    comb_init_flat = comb_init_flat * hc_scale[2] + hc_base[2 * hc_mult :].view(1, hc_mult, hc_mult)

    # 5. Initial Softmax over rows + first column normalization (align with original operator)
    comb_flat = comb_init_flat.clone()
    # Subtract row-wise max for numerical stability (avoid exp overflow)
    row_max = comb_flat.max(dim=-1, keepdim=True).values
    comb_flat = torch.exp(comb_flat - row_max)

    # 6. Sinkhorn iterations (alternating row/column normalization)
    for _ in range(sinkhorn_iters):
        # 6.1 Row normalization
        row_sum = comb_flat.sum(dim=-1, keepdim=True)
        comb_flat = comb_flat / (row_sum + eps)
        # 6.2 Column normalization
        col_sum = comb_flat.sum(dim=-2, keepdim=True)
        comb_flat = comb_flat / (col_sum + eps)

    # 7. Reshape back to original shape [b, s, ...]
    pre = pre_flat.view(b, s, hc_mult)
    post = post_flat.view(b, s, hc_mult)
    comb = comb_flat.view(b, s, hc_mult, hc_mult)

    return pre, post, comb

Model 接口

class SinkhornFunction(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        mixes: torch.Tensor,
        hc_scale: torch.Tensor,
        hc_base: torch.Tensor,
        hc_mult: int = 4,
        sinkhorn_iters: int = 20,
        eps: float = 1e-6,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Triton implementation of HC-Split Sinkhorn, optimized for GPU performance

        Args:
            mixes: Input tensor with shape [batch_size, seq_len, (2+hc_mult)*hc_mult]
            hc_scale: Scale tensor with shape [3] (pre/post/comb scales)
            hc_base: Base tensor with shape [(2+hc_mult)*hc_mult] (pre/post/comb bases)
            hc_mult: HC dimension size (only 4 supported in current implementation), default=4
            sinkhorn_iters: Number of Sinkhorn normalization iterations, default=20
            eps: Small constant to prevent division by zero, default=1e-6

        Returns:
            tuple: (pre, post, comb)
                - pre: Output tensor with shape [batch_size, seq_len, hc_mult]
                - post: Output tensor with shape [batch_size, seq_len, hc_mult]
                - comb: Output tensor with shape [batch_size, seq_len, hc_mult, hc_mult]
        """

输入说明

参数	类型	描述	支持的数据类型
mixes	torch.Tensor	[batch_size, seq_len, (2+hc_mult)*hc_mult]	float32
hc_scale	torch.Tensor	[3]	float32
hc_base	torch.Tensor	[(2+hc_mult)*hc_mult]	-
hc_mult	attr	立即数，当前仅支持4	int32
sinkhorn_iters	attr	立即数，当前仅支持20	-
eps	attr	立即数，默认为1e-6	-

输出说明

参数	类型	描述	支持的数据类型
pre	torch.Tensor	[batch_size, seq_len, hc_mult]	与输入张量相同的数据类型
post	torch.Tensor	[batch_size, seq_len, hc_mult]	与输入张量相同的数据类型
comb	torch.Tensor	[batch_size, seq_len, hc_mult, hc_mult]	与输入张量相同的数据类型

实现原理

算子通过以下步骤实现：

1. 逻辑任务计算：获取输入shape的b、s维做任务总数
2. 设置块大小：尾轴4在max指令下性能不佳，需要pad后使能向量化运算
3. 启动内核：调用 Triton 内核执行并行计算操作

使用示例

import torch
from mindspeed_ops.api.triton.sinkhorn import SinkhornFunction

# 创建测试张量
mixes = torch.randn((1, 2048, 24), requires_grad=True, dtype=dtype).to(device)
hc_scale = torch.randn((3,), requires_grad=True, dtype=dtype).to(device)
hc_base = torch.randn((24,), requires_grad=True, dtype=dtype).to(device)

# 执行加法操作
triton_pre, triton_post, triton_comb = SinkhornFunction.apply(triton_mixes, triton_hc_scale, triton_hc_base)

# 验证结果
print("sinkhorn operation exec successfully!")

性能对比

方法	时间消耗
PyTorch 内置逻辑	基准
Triton 实现逻辑	约 2.0x

注：性能提升数据基于昇腾 NPU 环境测试，具体数值可能因硬件配置不同而有所差异。

适用场景

算子适用于以下场景：

1. 科学计算：多个vec算子融合，该算子能够提供性能优势
2. 数据类型：fp16与bf16虽支持，但过程中涉及多次求和，建议使用fp32

注意事项

1. 输入要求详见参数说明
2. 对于非常小的张量，可能不会观察到明显的性能提升，因为启动内核的开销和Triton下发可能超过计算收益