FlagGems/src/flag_gems/fused/topk_softmax.py at 209b412b2a7bfc1265b2401f7f36da41ef52c49a · AdvancedCompiler/FlagGems · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import torch
import triton
import triton.language as tl


@triton.jit
def topk_gating_softmax_kernel(
    input_ptr,
    finished_ptr,  # interface reserved, not yet used
    output_ptr,
    indices_ptr,
    source_rows_ptr,
    num_rows,
    k,
    num_experts,
    start_expert,
    end_expert,
    renormalize,
    INDEX_TY: tl.constexpr,
    BLOCK_SIZE_ROWS: tl.constexpr,
    BLOCK_SIZE_EXPERTS: tl.constexpr,
):
    pid = tl.program_id(0)
    rows = tl.arange(0, BLOCK_SIZE_ROWS) + pid * BLOCK_SIZE_ROWS
    valid_rows = rows < num_rows

    cols = start_expert + tl.arange(0, BLOCK_SIZE_EXPERTS)
    valid_cols = cols < end_expert

    logits = tl.load(
        input_ptr + rows[:, None] * num_experts + cols[None, :],
        mask=valid_rows[:, None] & valid_cols[None, :],
        other=-float("inf"),
    ).to(tl.float32)

    row_max = tl.max(logits, axis=1)[:, None]
    exp_vals = tl.exp(logits - row_max)
    probs = exp_vals / (tl.sum(exp_vals, axis=1)[:, None] + 1e-8)

    selected_sum = tl.zeros([BLOCK_SIZE_ROWS], dtype=tl.float32)
    for ki in range(k):
        curr_max = tl.max(probs, axis=1)
        curr_arg = tl.argmax(probs, axis=1) + start_expert

        tl.store(output_ptr + rows * k + ki, curr_max, mask=valid_rows)
        tl.store(indices_ptr + rows * k + ki, curr_arg.to(INDEX_TY), mask=valid_rows)
        tl.store(
            source_rows_ptr + rows * k + ki,
            (ki * num_rows + rows).to(tl.int32),
            mask=valid_rows,
        )
        selected_sum += curr_max

        probs = tl.where(
            cols[None, :] == (curr_arg[:, None] - start_expert), -float("inf"), probs
        )

    if renormalize:
        norm = selected_sum + 1e-8
        for ki in range(k):
            idx = rows * k + ki
            val = tl.load(output_ptr + idx, mask=valid_rows)
            val = val / norm
            tl.store(output_ptr + idx, val, mask=valid_rows)


def topk_softmax(
    topk_weights: torch.Tensor,
    topk_indices: torch.Tensor,
    token_expert_indices: torch.Tensor,
    gating_output: torch.Tensor,
    renormalize: bool,
) -> None:
    num_tokens, num_experts = gating_output.shape
    topk = topk_weights.size(-1)
    assert topk <= 32

    if topk_indices.dtype == torch.int32:
        index_ty = tl.int32
    elif topk_indices.dtype == torch.uint32:
        index_ty = tl.uint32
    elif topk_indices.dtype == torch.int64:
        index_ty = tl.int64
    else:
        raise TypeError("topk_indices must be int32/int64/uint32")

    max_total_threads = 1024
    BLOCK_SIZE_EXPERTS = ((triton.next_power_of_2(num_experts) + 31) // 32) * 32
    BLOCK_SIZE_EXPERTS = min(BLOCK_SIZE_EXPERTS, 1024)
    BLOCK_SIZE_ROWS = max_total_threads // BLOCK_SIZE_EXPERTS
    BLOCK_SIZE_ROWS = max(BLOCK_SIZE_ROWS, 1)

    grid = (triton.cdiv(num_tokens, BLOCK_SIZE_ROWS),)

    topk_gating_softmax_kernel[grid](
        input_ptr=gating_output,
        finished_ptr=None,
        output_ptr=topk_weights,
        indices_ptr=topk_indices,
        source_rows_ptr=token_expert_indices,
        num_rows=num_tokens,
        k=topk,
        num_experts=num_experts,
        start_expert=0,
        end_expert=num_experts,
        renormalize=renormalize,
        INDEX_TY=index_ty,
        BLOCK_SIZE_ROWS=BLOCK_SIZE_ROWS,
        BLOCK_SIZE_EXPERTS=BLOCK_SIZE_EXPERTS,
    )