Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / linear_attention_convolution.py: 19%

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"""Optional NPU kernels for dense and packed depthwise causal convolution.""" 

16 

17from __future__ import annotations 

18 

19from typing import TYPE_CHECKING, Any 

20 

21import torch 

22from torch.nn import functional as F 

23 

24if TYPE_CHECKING: 

25 from hyper_parallel.components.functional.linear_attention_packed import PackedSequenceInfo 

26 

27 

28def validate_conv_backend(backend: str) -> str: 

29 """Validate an explicit convolution backend without loading kernels. 

30 

31 Args: 

32 backend: One of ``eager``, ``triton``, or ``ascendc``. 

33 

34 Returns: 

35 The validated backend, without changing its spelling or default. 

36 

37 Raises: 

38 ValueError: If the backend is unsupported. 

39 """ 

40 if backend not in ("eager", "triton", "ascendc"): 

41 raise ValueError(f"Unsupported convolution backend {backend!r}; expected eager, triton, or ascendc.") 

42 return backend 

43 

44 

45class _AscendCConvolution(torch.autograd.Function): 

46 """Shared dense/packed convolution with joint SiLU and convolution backward.""" 

47 

48 @staticmethod 

49 def forward( 

50 ctx: Any, 

51 inputs: torch.Tensor, 

52 weight: torch.Tensor, 

53 bias: torch.Tensor | None, 

54 offsets: tuple[int, ...] | None, 

55 ) -> torch.Tensor: 

56 """Save preactivation and per-call metadata for joint backward.""" 

57 # Optional native operators must not be loaded by ordinary framework imports. 

58 from fla_npu.ops.ascendc import causal_conv1d_fn # pylint: disable=import-outside-toplevel 

59 

60 op_inputs = inputs if offsets is None else inputs.reshape(-1, inputs.shape[-1]) 

61 op_inputs = op_inputs.contiguous() 

62 op_weight = weight[:, 0].transpose(0, 1).contiguous() 

63 cpu_offsets = None if offsets is None else torch.tensor(offsets, dtype=torch.int64, device="cpu") 

64 preactivation = causal_conv1d_fn( 

65 op_inputs, op_weight, bias, conv_states=None, activation=None, 

66 query_start_loc_cpu=cpu_offsets, head_num=0, 

67 ) 

68 ctx.save_for_backward(op_inputs, op_weight, preactivation, cpu_offsets) 

69 ctx.input_shape = tuple(inputs.shape) 

70 ctx.has_bias = bias is not None 

71 ctx.input_layout = "BSND" if offsets is None else "TND" 

72 return F.silu(preactivation).reshape_as(inputs) 

73 

74 @staticmethod 

75 def backward( 

76 ctx: Any, grad_output: torch.Tensor, 

77 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, None]: 

78 """Differentiate SiLU once inside the native convolution backward.""" 

79 # Keep the optional native dependency lazy in backward as well as forward. 

80 from fla_npu.ops.ascendc import causal_conv1d_bwd # pylint: disable=import-outside-toplevel 

81 

82 inputs, weight, preactivation, cpu_offsets = ctx.saved_tensors 

83 grad_input, grad_weight, grad_bias, _ = causal_conv1d_bwd( 

84 x=inputs, y=preactivation, weight=weight, 

85 dy=grad_output.reshape_as(inputs).contiguous(), 

86 query_start_loc=cpu_offsets, activation=1, input_layout=ctx.input_layout, 

87 ) 

88 return ( 

89 grad_input.reshape(ctx.input_shape), 

90 grad_weight.transpose(0, 1).unsqueeze(1).contiguous(), 

91 grad_bias if ctx.has_bias else None, 

92 None, 

93 ) 

94 

95 

96def _with_history( 

97 inputs: torch.Tensor, 

98 halo: torch.Tensor | None, 

99 info: PackedSequenceInfo | None, 

100 width: int, 

101) -> tuple[torch.Tensor, tuple[int, ...] | None, int]: 

102 """Prepend only the first fragment's true history without inventing a reset.""" 

103 if info is None: 

104 return inputs, None, 0 

105 if halo is None: 

106 return inputs, info.local_offsets, 0 

107 prefix = min(width - 1, info.preceding_sample_tokens) 

108 # The empty slice preserves a zero halo gradient when a real document resets. 

109 extended = torch.cat((halo[:, width - 1 - prefix:], inputs), dim=1) 

110 offsets = (0, *(offset + prefix for offset in info.local_offsets[1:])) 

111 return extended, offsets, prefix 

112 

113 

114def _triton_convolution( 

115 inputs: torch.Tensor, 

116 weight: torch.Tensor, 

117 bias: torch.Tensor | None, 

118 info: PackedSequenceInfo | None, 

119 offsets: tuple[int, ...] | None, 

120 use_plan: bool, 

121) -> torch.Tensor: 

122 """Use one FLA entry for dense batches and immutable packed document plans.""" 

123 # FLA is optional and must not be imported for the default eager backend. 

124 from fla.modules.conv import causal_conv1d # pylint: disable=import-outside-toplevel 

125 

126 if info is None: 

127 cu_seqlens, cu_seqlens_cpu, chunk_indices = None, None, None 

128 elif use_plan: 

129 plan = info.materialize(inputs.device, 64) 

130 cu_seqlens, cu_seqlens_cpu = plan.cu_seqlens, plan.cu_seqlens_cpu 

131 chunk_indices = plan.chunk_indices 

132 else: 

133 cu_seqlens_cpu = torch.tensor(offsets, dtype=torch.int64, device="cpu") 

134 cu_seqlens = cu_seqlens_cpu.to(device=inputs.device, dtype=torch.int32) 

135 chunk_indices = None 

136 # Contiguous BTD enables the NPU kernel without its strided fallback workspace. 

137 output, _ = causal_conv1d( 

138 inputs.contiguous(), weight=weight[:, 0], bias=bias, activation=None, backend="triton", 

139 cu_seqlens=cu_seqlens, cu_seqlens_cpu=cu_seqlens_cpu, chunk_indices=chunk_indices, 

140 ) 

141 return F.silu(output) 

142 

143 

144def _validate_fused_inputs( 

145 inputs: torch.Tensor, 

146 weight: torch.Tensor, 

147 bias: torch.Tensor | None, 

148 info: PackedSequenceInfo | None, 

149 halo: torch.Tensor | None, 

150) -> None: 

151 """Check layouts before a selected backend is allowed to launch kernels.""" 

152 if inputs.ndim != 3 or any(size < 1 for size in inputs.shape): 

153 raise ValueError("Fused convolution expects nonempty [B,S,C] or packed [1,T,C] inputs.") 

154 if weight.ndim != 3 or weight.shape[:2] != (inputs.shape[-1], 1): 

155 raise ValueError("Fused convolution weight must have shape [C,1,W] matching input channels.") 

156 if bias is not None and bias.shape != (inputs.shape[-1],): 

157 raise ValueError("Fused convolution bias must have shape [C] matching input channels.") 

158 if weight.device != inputs.device or (bias is not None and bias.device != inputs.device): 

159 raise ValueError("Fused convolution input, weight and bias must be on the same device.") 

160 if info is not None and inputs.shape[:2] != (1, info.length): 

161 raise ValueError("Packed convolution expects [1,T,C] matching its token interval.") 

162 if halo is not None: 

163 _validate_fused_history(inputs, halo, info, weight.shape[-1]) 

164 

165 

166def _validate_fused_history( 

167 inputs: torch.Tensor, 

168 halo: torch.Tensor, 

169 info: PackedSequenceInfo | None, 

170 width: int, 

171) -> None: 

172 """Reject history that cannot represent the same document and activation layout.""" 

173 if info is None: 

174 raise ValueError("Fused convolution halo requires explicit packed sequence metadata.") 

175 if halo.shape != (1, width - 1, inputs.shape[-1]): 

176 raise ValueError("Fused convolution halo has an incompatible shape.") 

177 if halo.dtype != inputs.dtype or halo.device != inputs.device: 

178 raise ValueError("Fused convolution halo must match the input dtype and device.") 

179 

180 

181def fused_causal_conv1d( 

182 inputs: torch.Tensor, 

183 weight: torch.Tensor, 

184 bias: torch.Tensor | None, 

185 info: PackedSequenceInfo | None = None, 

186 halo: torch.Tensor | None = None, 

187 *, 

188 backend: str, 

189 dilation: int = 1, 

190) -> torch.Tensor: 

191 """Apply fused depthwise causal convolution and SiLU to dense or packed tokens. 

192 

193 Args: 

194 inputs: Dense ``[B,S,C]`` or packed ``[1,T,C]`` projected tokens. 

195 weight: The original parameter's local view with shape ``[C,1,W]``. 

196 bias: Optional local bias with shape ``[C]``. 

197 info: Packed document boundaries, or None for independent dense batch samples. 

198 halo: Already exchanged packed history; this helper adds no communication. 

199 backend: Explicit ``triton`` or ``ascendc`` selection, with no fallback. 

200 dilation: Source convolution dilation; fused kernels require one. 

201 

202 Returns: 

203 Local SiLU output with the input shape and dtype. 

204 

205 Raises: 

206 ValueError: If the backend, device, dilation, or native dtype/shape is unsupported. 

207 ImportError: If the selected optional kernel package is missing. 

208 

209 Note: 

210 AscendC jointly differentiates SiLU and convolution using a saved low-precision 

211 preactivation. Unlike eager or Triton, its backward does not round the SiLU 

212 gradient to the input dtype before convolution backward. Dense batch samples 

213 remain separate; only packed inputs with explicit metadata are flattened. 

214 """ 

215 if backend not in ("triton", "ascendc"): 

216 raise ValueError(f"Fused convolution requires triton or ascendc, got {backend!r}.") 

217 if dilation != 1: 

218 raise ValueError("Fused convolution requires dilation=1.") 

219 if inputs.device.type != "npu": 

220 raise ValueError("Fused convolution currently supports NPU tensors only.") 

221 _validate_fused_inputs(inputs, weight, bias, info, halo) 

222 width = weight.shape[-1] 

223 if backend == "ascendc": 

224 if width not in (2, 3, 4) or inputs.shape[-1] % 16: 

225 raise ValueError("AscendC convolution requires width 2/3/4 and channels divisible by 16.") 

226 if inputs.dtype not in (torch.bfloat16, torch.float16): 

227 raise ValueError("AscendC convolution requires BF16 or FP16 activations.") 

228 if weight.dtype != inputs.dtype or (bias is not None and bias.dtype != inputs.dtype): 

229 raise ValueError("AscendC convolution requires matching input/weight/bias dtypes; no implicit cast.") 

230 extended, offsets, prefix = _with_history(inputs, halo, info, width) 

231 if backend == "triton": 

232 output = _triton_convolution(extended, weight, bias, info, offsets, halo is None) 

233 else: 

234 output = _AscendCConvolution.apply(extended, weight, bias, offsets) 

235 return output[:, prefix:]