Coverage for src/flag_gems/runtime/backend/_kunlunxin/ops/matmul_int8.py: 0%

1# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. 

2# 

3# Permission is hereby granted, free of charge, to any person obtaining a copy 

4# of this software and associated documentation files (the "Software"), to deal 

5# in the Software without restriction, including without limitation the rights 

6# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

7# copies of the Software, and to permit persons to whom the Software is 

8# furnished to do so, subject to the following conditions: 

9# 

10# The above copyright notice and this permission notice shall be included in 

11# all copies or substantial portions of the Software. 

12# 

13# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

14# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

15# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

16# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

17# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

18# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 

19# THE SOFTWARE. 

20 

21""" 

22Matrix Multiplication 

23=============== 

24""" 

25 

26import torch 

27import triton 

28import triton.language as tl 

29 

30DEV = "cuda" 

31 

32 

33def get_output_dtype(a_dtype, b_dtype): 

34 return torch.bfloat16 

35 

36 

37def get_autotune_config(): 

38 return [ 

39 triton.Config({"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}), 

40 triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 128}), 

41 triton.Config({"BLOCK_SIZE_M": 256, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 256}), 

42 ] 

43 

44 

45@triton.autotune( 

46 configs=get_autotune_config(), 

47 key=["M", "N", "K"], 

48) 

49@triton.jit 

50def matmul_kernel( 

51 # Pointers to matrices 

52 a_ptr, 

53 b_ptr, 

54 c_ptr, 

55 # Matrix dimensions 

56 M, 

57 N, 

58 K, 

59 # The stride variables represent how much to increase the ptr by when moving by 1 

60 # element in a particular dimension. 

61 stride_am, 

62 stride_ak, # 

63 stride_bk, 

64 stride_bn, # 

65 stride_cm, 

66 stride_cn, 

67 # Meta-parameters 

68 BLOCK_SIZE_M: tl.constexpr, 

69 BLOCK_SIZE_N: tl.constexpr, 

70 BLOCK_SIZE_K: tl.constexpr, # 

71): 

72 """Kernel for computing the matmul C = A x B. 

73 A has shape (M, K), B has shape (K, N) and C has shape (M, N) 

74 """ 

75 # L2 Cache Optimization: Group multiple M-blocks together to reuse B columns 

76 # GROUP_SIZE_M=8 means 8 consecutive M-blocks share the same B columns in L2 cache 

77 GROUP_SIZE_M: tl.constexpr = 8 

78 # ----------------------------------------------------------- 

79 # Map program ids `pid` to the block of C it should compute. 

80 # This is done in a grouped ordering to promote L2 data reuse. 

81 # See above `L2 Cache Optimizations` section for details. 

82 pid = tl.program_id(axis=0) 

83 num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) 

84 num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) 

85 num_pid_in_group = GROUP_SIZE_M * num_pid_n 

86 group_id = pid // num_pid_in_group 

87 first_pid_m = group_id * GROUP_SIZE_M 

88 group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) 

89 pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) 

90 pid_n = (pid % num_pid_in_group) // group_size_m 

91 

92 # ---------------------------------------------------------- 

93 # Create block pointers for A, B, and C using make_block_ptr. 

94 a_block_ptr = tl.make_block_ptr( 

95 base=a_ptr, 

96 shape=(M, K), 

97 strides=(stride_am, stride_ak), 

98 offsets=(pid_m * BLOCK_SIZE_M, 0), 

99 block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K), 

100 order=(1, 0), 

101 ) 

102 b_block_ptr = tl.make_block_ptr( 

103 base=b_ptr, 

104 shape=(K, N), 

105 strides=(stride_bk, stride_bn), 

106 offsets=(0, pid_n * BLOCK_SIZE_N), 

107 block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N), 

108 order=(1, 0), 

109 ) 

110 # ----------------------------------------------------------- 

111 # Iterate to compute a block of the C matrix. 

112 # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block 

113 # of fp32 values for higher accuracy. 

114 accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) 

115 for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): 

116 a = tl.load(a_block_ptr, boundary_check=(0, 1), padding_option="zero") 

117 b = tl.load(b_block_ptr, boundary_check=(0, 1), padding_option="zero") 

118 accumulator = tl.dot(a.to(tl.bfloat16), b.to(tl.bfloat16), accumulator) 

119 a_block_ptr = tl.advance(a_block_ptr, (0, BLOCK_SIZE_K)) 

120 b_block_ptr = tl.advance(b_block_ptr, (BLOCK_SIZE_K, 0)) 

121 c = accumulator.to(c_ptr.dtype.element_ty) 

122 # ----------------------------------------------------------- 

123 # Write back the block of the output matrix C. 

124 c_block_ptr = tl.make_block_ptr( 

125 base=c_ptr, 

126 shape=(M, N), 

127 strides=(stride_cm, stride_cn), 

128 offsets=(pid_m * BLOCK_SIZE_M, pid_n * BLOCK_SIZE_N), 

129 block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N), 

130 order=(1, 0), 

131 ) 

132 tl.store(c_block_ptr, c, boundary_check=(0, 1)) 

133 

134 

135def torch_matmul(a, b): 

136 print(f"{a.dtype=} {b.dtype=}") 

137 # b is (N, K), so b.t() gives (K, N) 

138 c = torch.matmul(a.to(torch.bfloat16), b.to(torch.bfloat16).t()) 

139 return c 

140 

141 

142# %% 

143# We can now create a convenience wrapper function that only takes two input tensors, 

144# and (1) checks any shape constraint; (2) allocates the output; (3) launches the above kernel. 

145 

146 

147def matmul_int8(a, b): 

148 # Save original shape for 3D support 

149 a_shape = a.shape 

150 if a.ndim == 3: 

151 a = a.contiguous().reshape(-1, a.shape[-1]) 

152 # Handle non-contiguous inputs if necessary 

153 if a.stride(0) > 1 and a.stride(1) > 1: 

154 a = a.contiguous() 

155 # b has shape (N, K), transpose to (K, N) contiguous for the kernel 

156 b = b.t().contiguous() 

157 # Check constraints. After transpose, b has shape (K, N) 

158 assert a.shape[1] == b.shape[0], "Incompatible dimensions" 

159 M, K = a.shape 

160 N = b.shape[1] 

161 # Allocates output. 

162 c_dtype = get_output_dtype(a.dtype, b.dtype) 

163 c = torch.empty((M, N), device=a.device, dtype=c_dtype) 

164 # 1D launch kernel where each block gets its own program. 

165 grid = lambda META: ( 

166 triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), 

167 ) 

168 matmul_kernel[grid]( 

169 a, 

170 b, 

171 c, # 

172 M, 

173 N, 

174 K, # 

175 a.stride(0), 

176 a.stride(1), # 

177 b.stride(0), 

178 b.stride(1), 

179 c.stride(0), 

180 c.stride(1), # 

181 ) 

182 # Reshape output back if input was 3D 

183 if len(a_shape) == 3: 

184 c = c.reshape(*a_shape[:-1], N) 

185 return c