cera 0.5.5

Rust-native LLM inference engine
Documentation 
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
111
112
113
114
115
116
117
118
119
120
121
122
123

#include <metal_stdlib>
#include <metal_math>
#include <metal_texture>
using namespace metal;

#line 60 "cera/src/backend/shaders/slang/moe_combine.slang"
struct KernelContext_0
{
    packed_uint4 device* params_0;
    float device* sel_weight_0;
    float device* z_0;
    float device* out_buf_0;
};


#line 42
[[kernel]] void moe_combine(uint3 lid_0 [[thread_position_in_threadgroup]], uint3 grp_0 [[threadgroup_position_in_grid]], packed_uint4 device* params_1 [[buffer(3)]], float device* sel_weight_1 [[buffer(1)]], float device* z_1 [[buffer(0)]], float device* out_buf_1 [[buffer(2)]])
{

#line 42
    thread KernelContext_0 kernelContext_0;

#line 42
    (&kernelContext_0)->params_0 = params_1;

#line 42
    (&kernelContext_0)->sel_weight_0 = sel_weight_1;

#line 42
    (&kernelContext_0)->z_0 = z_1;

#line 42
    (&kernelContext_0)->out_buf_0 = out_buf_1;
    uint hidden_0 = (uint4(*(params_1+int(0))) ).x;
    uint _S1 = (uint4(*(params_1+int(0))) ).y;
    uint n_tokens_0 = (uint4(*(params_1+int(0))) ).z;
    bool accumulate_0 = ((uint4(*(params_1+int(0))) ).w) != 0U;

    uint row_0 = grp_0.x * 256U + lid_0.x;
    uint tok_0 = grp_0.y;

#line 49
    bool _S2;
    if(row_0 >= hidden_0)
    {

#line 50
        _S2 = true;

#line 50
    }
    else
    {

#line 50
        _S2 = tok_0 >= n_tokens_0;

#line 50
    }

#line 50
    if(_S2)
    {

#line 51
        return;
    }

#line 51
    uint s_0 = 0U;

#line 51
    float acc_0 = 0.0f;



    for(;;)
    {

#line 55
        if(s_0 < _S1)
        {
        }
        else
        {

#line 55
            break;
        }

#line 56
        uint entry_0 = tok_0 * _S1 + s_0;
        float acc_1 = acc_0 + (&kernelContext_0)->sel_weight_0[entry_0] * (&kernelContext_0)->z_0[entry_0 * hidden_0 + row_0];

#line 55
        s_0 = s_0 + 1U;

#line 55
        acc_0 = acc_1;

#line 55
    }



    uint dst_0 = tok_0 * hidden_0 + row_0;
    float device* _S3 = (&kernelContext_0)->out_buf_0+dst_0;

#line 60
    if(accumulate_0)
    {

#line 60
        acc_0 = *((&kernelContext_0)->out_buf_0+dst_0) + acc_0;

#line 60
    }

#line 60
    *_S3 = acc_0;
    return;
}