Kernels:

kernels-community
/

fla

Trusted publisher

Uploaded using `kernel-builder`.

e0d830e verified 12 days ago

4.06 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):
    can_torch_compile = True

    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):
    can_torch_compile = True

    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):
    can_torch_compile = True

    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):
    can_torch_compile = True

    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):
    can_torch_compile = True

    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"
]

	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):
	can_torch_compile = True

	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):
	can_torch_compile = True

	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):
	can_torch_compile = True

	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):
	can_torch_compile = True

	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):
	can_torch_compile = True

	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"
	]