Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / custom_ops / gdn / triton / cumsum.py: 0%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15# -*- coding: utf-8 -*- 

16# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang 

17 

18# pylint: disable=missing-public-type-hints,missing-public-docstring,disallowed-name 

19# pylint: disable=useless-return,unused-argument,no-else-return,invalid-name 

20# pylint: disable=missing-module-docstring,missing-function-docstring 

21 

22from typing import Optional 

23 

24import torch 

25import triton 

26import triton.language as tl 

27 

28from .utils import prepare_chunk_indices 

29 

30 

31@triton.heuristics({ 

32 'HAS_SCALE': lambda args: args['scale'] is not None, 

33 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None 

34}) 

35@triton.jit(do_not_specialize=['T']) 

36def chunk_local_cumsum_scalar_kernel( 

37 s, 

38 o, 

39 scale, 

40 cu_seqlens, 

41 chunk_indices, 

42 T, 

43 B: tl.constexpr, 

44 H: tl.constexpr, 

45 BLOCK_T: tl.constexpr, 

46 REVERSE: tl.constexpr, 

47 HAS_SCALE: tl.constexpr, 

48 IS_VARLEN: tl.constexpr, 

49 HEAD_FIRST: tl.constexpr, 

50 CHUNK_SIZE: tl.constexpr = 64, 

51): 

52 i_block, i_b = tl.program_id(0), tl.program_id(1) 

53 N_CHUNKS: tl.constexpr = BLOCK_T // CHUNK_SIZE 

54 

55 if IS_VARLEN: 

56 i_s, i_block = tl.load(chunk_indices + i_block * 2).to(tl.int32), tl.load( 

57 chunk_indices + i_block * 2 + 1 

58 ).to(tl.int32) 

59 

60 bos, eos = tl.load(cu_seqlens + i_s).to(tl.int32), tl.load( 

61 cu_seqlens + i_s + 1 

62 ).to(tl.int32) 

63 T = eos - bos 

64 else: 

65 bos, eos = i_b * T, i_b * T + T 

66 

67 ptr_s = tl.make_block_ptr( 

68 s + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0) 

69 ) 

70 ptr_o = tl.make_block_ptr( 

71 o + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0) 

72 ) 

73 b_s = tl.load(ptr_s, boundary_check=(0,)).to(tl.float32) 

74 b_s = tl.reshape(b_s, (N_CHUNKS, CHUNK_SIZE, H)) 

75 b_s = tl.trans(b_s, (1, 0, 2)) 

76 b_o = tl.cumsum(b_s, axis=0) 

77 if REVERSE: 

78 b_z = tl.sum(b_s, axis=0) 

79 b_o = -b_o + b_z[None] + b_s 

80 if HAS_SCALE: 

81 b_o *= scale 

82 b_o = tl.trans(b_o, (1, 0, 2)) 

83 b_o = tl.reshape(b_o, (BLOCK_T, H)) 

84 

85 tl.store(ptr_o, b_o.to(ptr_o.dtype.element_ty), boundary_check=(0,)) 

86 return 

87 

88 

89def chunk_local_cumsum_scalar( 

90 g: torch.Tensor, 

91 chunk_size: int, 

92 reverse: bool = False, 

93 scale: float = None, 

94 cu_seqlens: Optional[torch.Tensor] = None, 

95 head_first: bool = False, 

96 output_dtype: Optional[torch.dtype] = torch.float 

97) -> torch.Tensor: 

98 

99 B, T, H = g.shape 

100 if chunk_size != 2 ** (chunk_size.bit_length() - 1): 

101 raise ValueError( 

102 f"chunk_size must be a power of 2, chunk_size is {chunk_size}" 

103 ) 

104 # We adjust the tiling strategy to prevent overflow in in backward passes and context parallel scenarios 

105 # while maximizing UB utilization where possible. 

106 # The tiling strategy is as follows: 

107 # 1. BT must be greater than or equal to chunk_size. 

108 # 2. UB estimation varies directly with H. 

109 # 3. BT in reverse mode is smaller than in forward mode. 

110 BT = max(chunk_size, triton.next_power_of_2((1 << 11 if reverse else 1 << 12) // H)) 

111 chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None 

112 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) 

113 g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) 

114 grid = (NT, B) 

115 chunk_local_cumsum_scalar_kernel[grid]( 

116 s=g_org, 

117 o=g, 

118 scale=scale, 

119 cu_seqlens=cu_seqlens, 

120 chunk_indices=chunk_indices, 

121 T=T, 

122 B=B, 

123 H=H, 

124 BLOCK_T=BT, 

125 HEAD_FIRST=head_first, 

126 REVERSE=reverse, 

127 CHUNK_SIZE=chunk_size, 

128 ) 

129 return g 

130 

131 

132def chunk_local_cumsum( 

133 g: torch.Tensor, 

134 chunk_size: int, 

135 reverse: bool = False, 

136 scale: float = None, 

137 cu_seqlens: Optional[torch.Tensor] = None, 

138 head_first: bool = False, 

139 output_dtype: Optional[torch.dtype] = torch.float, 

140 **kwargs 

141) -> torch.Tensor: 

142 if cu_seqlens is not None: 

143 if g.shape[0] != 1: 

144 raise ValueError( 

145 "Only batch size 1 is supported when cu_seqlens are provided, " 

146 f"current size is {g.shape[0]}" 

147 ) 

148 if len(g.shape) == 3: 

149 return chunk_local_cumsum_scalar( 

150 g=g, 

151 chunk_size=chunk_size, 

152 reverse=reverse, 

153 scale=scale, 

154 cu_seqlens=cu_seqlens, 

155 head_first=head_first, 

156 output_dtype=output_dtype 

157 ) 

158 else: 

159 raise ValueError( 

160 f"Unsupported input shape {g.shape}, " 

161 f"which should be (B, T, H, D) if `head_first=False` " 

162 f"or (B, H, T, D) otherwise" 

163 )