Kernels:

kernels-community
/

fla

Trusted publisher

Uploaded using `kernel-builder`.

888b93f verified 1 day ago

3.91 kB

	import torch.nn as nn

	from .modules.fused_norm_gate import rms_norm_gated
	from .ops.gated_delta_rule import chunk_gated_delta_rule as chunk_gdn
	from .ops.gated_delta_rule import fused_recurrent_gated_delta_rule as fused_recurrent_gdn
	from .ops.kda import chunk_kda, fused_recurrent_kda


	class FusedRMSNormGated(nn.Module):
	def forward(self, hidden_states, gate=None):
	return rms_norm_gated(
	hidden_states,
	gate,
	self.weight,
	None, # bias
	self.activation,
	residual=None,
	eps=self.variance_epsilon,
	prenorm=False,
	residual_in_fp32=False,
	)


	class chunk_kimi_delta_attention(nn.Module):
	def forward(
	self,
	query,
	key,
	value,
	g,
	beta,
	chunk_size=64,
	initial_state=None,
	output_final_state=False,
	use_qk_l2norm_in_kernel=False,
	**kwargs,
	):
	# Keep internal consistency between transformers and fla to allow both
	cu_seqlens = kwargs.pop("cu_seq_lens_q", kwargs.pop("cu_seqlens", None))

	return chunk_kda(
	query,
	key,
	value,
	g=g,
	beta=beta,
	chunk_size=chunk_size,
	initial_state=initial_state,
	output_final_state=output_final_state,
	use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
	cu_seqlens=cu_seqlens,
	**kwargs,
	)


	class recurrent_kimi_delta_attention(nn.Module):
	def forward(
	self, query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False, **kwargs,
	):
	# Keep internal consistency between transformers and fla to allow both
	cu_seqlens = kwargs.pop("cu_seq_lens_q", kwargs.pop("cu_seqlens", None))

	return fused_recurrent_kda(
	query,
	key,
	value,
	g=g,
	beta=beta,
	initial_state=initial_state,
	output_final_state=output_final_state,
	use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
	cu_seqlens=cu_seqlens,
	**kwargs,
	)


	class chunk_gated_delta_rule(nn.Module):
	def forward(
	self,
	query,
	key,
	value,
	g,
	beta,
	chunk_size=64,
	initial_state=None,
	output_final_state=False,
	use_qk_l2norm_in_kernel=False,
	**kwargs,
	):
	# Keep internal consistency between transformers and fla to allow both
	cu_seqlens = kwargs.pop("cu_seq_lens_q", kwargs.pop("cu_seqlens", None))

	return chunk_gdn(
	query,
	key,
	value,
	g=g,
	beta=beta,
	chunk_size=chunk_size,
	initial_state=initial_state,
	output_final_state=output_final_state,
	use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
	cu_seqlens=cu_seqlens,
	**kwargs,
	)


	class recurrent_gated_delta_rule(nn.Module):
	def forward(
	self, query, key, value, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel=False, **kwargs,
	):
	# Keep internal consistency between transformers and fla to allow both
	cu_seqlens = kwargs.pop("cu_seq_lens_q", kwargs.pop("cu_seqlens", None))

	return fused_recurrent_gdn(
	query,
	key,
	value,
	g=g,
	beta=beta,
	initial_state=initial_state,
	output_final_state=output_final_state,
	use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
	cu_seqlens=cu_seqlens,
	**kwargs,
	)


	__all__ = [
	"FusedRMSNormGated",
	"chunk_kimi_delta_attention",
	"recurrent_kimi_delta_attention",
	"chunk_gated_delta_rule",
	"recurrent_gated_delta_rule"
	]