
    )`i              	       v    d dl Z d dlmZ d dlmZ e j        dej        dej        dej        ddfd            ZdS )    N)scale_and_clamp
BLOCK_SIZEHAS_X_SCALEHAS_O_SCALEreturnc
                    t          j        d                              t           j                  }
t          j        d          }| ||
z  z   }|||
z  z   }||z  t          j        d|          z   }||k     }t          j        ||z   |                              t           j                  }t          j        ||z   |z   |                              t           j                  }|rt          j        |          }||z  }||z  }t          j        |          |z  |z  }|	r/t          j        |          }t          ||| j	        j
                  }t          j        ||z   ||           dS )a  Sigmoid Linear Unit and Multiplication Kernel

    Args:
        o_ptr:       Pointer to the 2D output tensor.
        o_stride:    Output tensor stride.
        o_scale_ptr: The optional, known scale of the output activations.
        x_ptr:       Pointer to the 2D input tensor.
        x_stride:    Input tensor stride.
        x_scale_ptr: The optional, known scale of the input tensor.
        d:           The number of elements along the second dimension.
        BLOCK_SIZE:  Tunable block size to process in each kernel.

    Operating on a 2D grid, computes the following:

    ```
    out[i, j] = sigmoid(x[i, j]) * x[i, j] * x[i, j + d]
    ```

    If scales are provided, the input and output tensors are scaled.
    r   )axis   )maskN)tl
program_idtoint64arangeloadfloat32sigmoidr   dtype
element_tystore)o_ptro_strideo_scale_ptrx_ptrx_stridex_scale_ptrdr   r   r   ij	o_row_ptr	x_row_ptroffsetsr   abx_scaleresulto_scales                        x/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/flashinfer/triton/kernels/activation.pysilu_and_mul_kernelr)      s_   D 	1  **A
1A1$I1$I*nryJ777GQ;D
	G#$///222:>>A
	G#a'd33366rzBBA '+&&	W	WZ]]Q"F J'+&& %+2HIIHY &t444444    )	tritontriton.languagelanguager   flashinfer.triton.kernels.quantr   jit	constexprr)    r*   r(   <module>r2      s           ; ; ; ; ; ; 85 85 85 85 
85 85 85 85 85 85r*   