`i~ ddlZddlmZmZddlZddlmZddlmZm Z m Z m Z m Z ddl mZddlmZddlmZGddZGd d Zd ed ejd ejdefdZdedejdefdZdejdeeeefdedeeeeeefffdZGddejZGddZdS)N)TupleOptional)Boolean)Int32Float32minextract_mlir_valuesnew_from_mlir_values) HardwareInfo) WorkTileInfoc@eZdZdZddddedejfdZdZdZ dS) FmhaStaticTileSchedulerParamsaA class to represent parameters for the FMHA (Fused Multi-Head Attention) static tile scheduler. This class holds the configuration parameters needed to initialize and configure the tile scheduler for FMHA operations. :ivar is_persistent: Whether to use persistent kernel mode. :type is_persistent: bool :ivar problem_shape_mbh: Problem shape in (M, B, H) format. :type problem_shape_mbh: cute.Shape Nlocip is_persistentproblem_shape_mbhc>||_||_||_||_dS)a3 Initializes the FmhaStaticTileSchedulerParams with the given parameters. :param is_persistent: Whether to use persistent kernel mode. :type is_persistent: bool :param problem_shape_mbh: Problem shape in (M, B, H) format. :type problem_shape_mbh: cute.Shape N)rr_loc_ip)selfrrrrs /home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/cudnn/native_sparse_attention/compression/fmha_helpers.py__init__z&FmhaStaticTileSchedulerParams.__init__-s& +!2 cggc}|_|jfD]=}t|}||z }|jt |>|SN) _values_posrr appendlen)rvaluesobj obj_valuess r__extract_mlir_values__z5FmhaStaticTileSchedulerParams.__extract_mlir_values__Bsb#%r  *+ 5 5C,S11J j F   # #C OO 4 4 4 4 rc g}t|jg|jD]:\}}|t ||d|||d};t |jgt|Rd|jiS)Nr) ziprrrr rrtupler)rr obj_listr!n_itemss r__new_from_mlir_values__z6FmhaStaticTileSchedulerParams.__new_from_mlir_values__Js!7 8$:JKK & &LC OO0fXgX6FGG H H HGHH%FF,T-?c5??cccY]Ybcccr) __name__ __module__ __qualname____doc__boolcuteShaperr#r)rrrr!s{      :*dddddrrc eZdZdZddddededejdejfdZ e dddded ejfd Z e d e d ed ed e fdZdddd efdZddddZdddddZdZdZdS)FmhaStaticTileScheduleraiA static tile scheduler for FMHA (Fused Multi-Head Attention) operations. This class manages the scheduling of work tiles for FMHA kernels, supporting both persistent and non-persistent kernel modes. It tracks the current work position and advances through the problem space efficiently. :ivar _params: Scheduler parameters. :type _params: FmhaStaticTileSchedulerParams :ivar _blk_coord: Block coordinates. :type _blk_coord: cute.Coord :ivar _grid_shape: Grid shape for the kernel. :type _grid_shape: cute.Shape :ivar _is_persistent: Whether to use persistent kernel mode. :type _is_persistent: bool :ivar _current_work_linear_idx: Current linear work index. :type _current_work_linear_idx: Int32 :ivar _problem_shape_mbh: Problem shape in (M, B, H) format. :type _problem_shape_mbh: cute.Layout :ivar _num_blocks: Number of blocks in the problem. :type _num_blocks: Int32 :ivar _is_first_block: Whether this is the first block. :type _is_first_block: bool :ivar num_persistent_sm: Number of persistent SMs. :type num_persistent_sm: Int32 Nrparamscurrent_work_linear_idx blk_coord grid_shapec<||_||_||_|j|_||_t j|j|||_ t j |j |||_ d|_ t j ||||_ ||_||_dS)a Initializes the FmhaStaticTileScheduler with the given parameters. :param params: Scheduler parameters. :type params: FmhaStaticTileSchedulerParams :param current_work_linear_idx: Current linear work index. :type current_work_linear_idx: Int32 :param blk_coord: Block coordinates. :type blk_coord: cute.Coord :param grid_shape: Grid shape for the kernel. :type grid_shape: cute.Shape rTN)_params _blk_coord _grid_shaper_is_persistent_current_work_linear_idxr/ make_layoutr_problem_shape_mbhsize _num_blocks_is_first_blocknum_persistent_smrr)rr4r5r6r7rrs rrz FmhaStaticTileScheduler.__init__ms, #%$2(?%"&"263KQTY["\"\"\9T%<#"MMM#!%:32!F!F!F rreturnc|jrOt}|}t|t j|j||ddfS|jS)a Determine the grid shape for the FMHA kernel. For persistent kernels, the grid shape is limited by the number of SMs (Streaming Multiprocessors) available on the device. For non-persistent kernels, the grid shape matches the problem shape. :param params: Scheduler parameters. :type params: FmhaStaticTileSchedulerParams :return: Grid shape as (M, B, H) tuple. :rtype: cute.Shape r)rr get_device_multiprocessor_countrr/r@r)r4rr hardware_infosm_counts rget_grid_shapez&FmhaStaticTileScheduler.get_grid_shapese(   ,(NNM$DDFFHHdi(@cbQQQRR  + +rq_tiler current_idxseqlen_qc||z|kS)a7 Check if the current work index is valid for the given query sequence length. This method verifies that the current work tile index multiplied by the query tiler size is within the bounds of the query sequence length. :param q_tiler: Query tiler size. :type q_tiler: int :param current_idx: Current work index. :type current_idx: Int32 :param seqlen_q: Query sequence length. :type seqlen_q: Int32 :return: True if the work is valid, False otherwise. :rtype: Boolean r1)rKrLrMs rcheck_valid_work_for_seqlen_qz5FmhaStaticTileScheduler.check_valid_work_for_seqlen_qs,W$x//rc|jr|j|jkn|j}d}|jr#|j|j||}n|j}|dd|d|dff}t||S)aA Get information about the current work tile. Determines if the current work is valid and computes the tile coordinates based on whether the kernel is persistent or non-persistent. :return: WorkTileInfo containing tile coordinates and validity flag. :rtype: WorkTileInfo )rrrrrrF)r<r=rArBr?get_hier_coordr:r )rrris_validr6cur_tile_coords rget_current_workz(FmhaStaticTileScheduler.get_current_worksHLGZt4043CCC`d`t   (/>>t?\bejl>mmIII aL q\9Q< (  NH555rc0|||S)z} Get the initial work tile information. :return: Initial WorkTileInfo. :rtype: WorkTileInfo r)rU)rrrs rinitial_work_tile_infoz.FmhaStaticTileScheduler.initial_work_tile_infos$$$444rrF) advance_countrrcR|jr|xj||jzz c_d|_dS)a7 Advance to the next work tile. For persistent kernels, advances by the number of persistent SMs. For non-persistent kernels, marks that the first block has been processed. :param advance_count: Number of steps to advance (default: 1). :type advance_count: int FN)r<r=rCrB)rrXrrs radvance_to_next_workz,FmhaStaticTileScheduler.advance_to_next_works:   T  ) )]T=S-S S ) )$rct|j}|t|j|t|j|t|j|Sr)r r9extendr=r:r;)rr s rr#z/FmhaStaticTileScheduler.__extract_mlir_values__so$T\22 )$*GHHIII )$/::;;; )$*:;;<<< rc6t|dksJt|j|dd}t|j|dg}t|j|dd}t|j|dd}t ||||S)N r)rr r9r=r:r;r3)rr new_paramsnew_current_work_linear_idx new_blk_coordnew_grid_shapes rr)z0FmhaStaticTileScheduler.__new_from_mlir_values__s6{{b    )$,qs DD &:4;X[abc[dZe&f&f#,T_fQqSkJJ -d.>qrr KK&z3NP]_mnnnr)r*r+r,r-rrr/Coordr0r staticmethodrJintrrOr rUrWrZr#r)r1rrr3r3RsB     - "' :  J    F   ,,,-,  ,,,\,<0000  000\0.'+t666 66666-1T55555564D % % % % %ooooorr3r4r6r7rDc2t||d||S)aq Create a new FMHA static tile scheduler. :param params: Scheduler parameters. :type params: FmhaStaticTileSchedulerParams :param blk_coord: Block coordinates. :type blk_coord: cute.Coord :param grid_shape: Grid shape. :type grid_shape: cute.Shape :return: New FmhaStaticTileScheduler instance. :rtype: FmhaStaticTileScheduler r)r3)r4r6r7s r!create_fmha_static_tile_schedulerrj s$ #69Q<J O OOrrrc"t||S)ad Create FMHA static tile scheduler parameters. :param is_persistent: Whether to use persistent kernel mode. :type is_persistent: bool :param problem_shape_mbh: Problem shape in (M, B, H) format. :type problem_shape_mbh: cute.Shape :return: New FmhaStaticTileSchedulerParams instance. :rtype: FmhaStaticTileSchedulerParams )r)rrs r(create_fmha_static_tile_scheduler_paramsrls )8I J JJro_shape cta_tilerc >t|tjtj|d|dtj|ddtj|ddf}t|}||fS)ak Compute grid parameters for FMHA operation. This function calculates the appropriate grid shape and scheduler parameters based on the output tensor shape, CTA (Cooperative Thread Array) tiler, and whether to use persistent kernel mode. The output tensor o has shape (s, d, ((h_r, h_k), b)) where: - s: sequence length - d: head dimension - h_r: number of heads for query - h_k: number of heads for key - b: batch size :param o_shape: Output tensor shape for grid computation. :type o_shape: cute.Shape :param cta_tiler: CTA tiler dimensions (M, N, K). :type cta_tiler: Tuple[int, int, int] :param is_persistent: Whether to use persistent kernel mode. :type is_persistent: bool :return: Tuple of (scheduler_params, grid_shape). :rtype: Tuple[FmhaStaticTileSchedulerParams, Tuple[int, int, int]] rrQrF)rlr/ceil_divr@r3rJ)rmrnrtile_sched_paramsgrids r compute_gridrs0s:A M$)GAJ//1 > > Igajm $ $ Igajm $ $  # 1 12C D DD d ""rceZdZdZejZejZejZejZ dS)MaskTypeaNEnumeration of mask types for FMHA operations. - RESIDUAL_MASK: Residual mask for handling variable sequence lengths - WINDOW_MASK: Window mask for attention which also includes causal and no mask - WINDOW_MASK_INFERENCE: Same as the window mask, but has the limitation that the end of q is aligned with the end of k N) r*r+r,r-enumauto RESIDUAL_MASK WINDOW_MASKWINDOW_MASK_INFERENCECOMPRESSED_CAUSAL_MASKr1rrruru_sRDIKKM$)++K%DIKK&TY[[rruceZdZdZej ddedejdejde de de e d e e d e fd Z ej ddedejdejde de de e d e fd Z ej ddedejdejde de de e d e e d e e e ffd Zej ddedejdejde de de e d e e d e e e ffdZej ddedejdejde de de e d e e d e fdZej ddedejdejde de de e d e e de e d e fdZej ddedejdejde de de e d e e d e fdZej ddedejdejde de de ed e efdZdS) FusedMaskaA fused mask implementation for FMHA operations. This class handles different types of attention masks including no mask, residual mask for variable sequence lengths, and causal mask for autoregressive attention patterns. The class provides methods to: - Calculate trip counts for different mask types - Apply masks to attention scores - Handle masked and unmasked trip calculations N mask_typer6 tile_shaperMseqlen_kwindow_size_leftwindow_size_rightrDcd}tj|tjurdn||z }tj|tjkrt j||d}tj|tjkp|tjkrtj|durt j||d}n|ddz|dz} | |z|z} t j| |d} t j||d} t| | }ntj|tj kr|||z} |ddz|dzdz |z|z}t j|dz| z|d} t j||d} tdt| | }t ||||||}||z }|S)a Calculate the number of trips needed for the current block. The trip count depends on the mask type and the block coordinates. For causal masks, it considers the autoregressive constraint. :param mask_type: Type of mask to use :type mask_type: utils.MaskType :param blk_coord: Block coordinates. :type blk_coord: cute.Coord :param tile_shape: Shape of the tile. :type tile_shape: cute.Shape :param seqlen_q: Query sequence length for attention computation. :type seqlen_q: Int32 :param seqlen_k: Key sequence length for attention computation. :type seqlen_k: Int32 :param window_size_left: Left-side sliding window size for attention masking. :type window_size_left: Optional[Int32] :param window_size_right: Right-side sliding window size for attention masking. :type window_size_right: Optional[Int32] :return: Number of trips needed. :rtype: Int32 rrFN) cutlass const_exprrurzrxr/rpryrr{maxr}get_trip_start)r~r6rrMrrrresultoffset max_idx_qidx_k tmp_blocks_k max_blocks_kcompression_factor block_end start_blocks rget_trip_countzFusedMask.get_trip_countzsD((:X)XYYr_gjr_r  i8+AA B B <]8Z];;F  i8+??n9PXPnCn o o =!"3t";<< 9xA??&q\A-A> !F*->>#}UJqMBB #}Xz!}EE \<88   X-L L M M =!)X!5 "1)Z]:Q>GJ[[I=9q==O*OR\]^R_``L=:a=AALC l;;<((:X)XYYr_gjr_r&55iJX`bjl|}}--         q\A%A6?BRR s9 1 #=qAA:PbCbefCfgg!#333rc ntj|tjurdn||z }t||||||}t|||||||} |d|dz|z|z} tt| |dz| |zdz d} | |zdz } | | fS)a Get the begin and end tile idx for the trailing mask. :param mask_type: Type of mask to use :type mask_type: utils.MaskType :param blk_coord: Block coordinates. :type blk_coord: cute.Coord :param tile_shape: Shape of the tile. :type tile_shape: cute.Shape :param seqlen_q: Query sequence length for attention computation. :type seqlen_q: Int32 :param seqlen_k: Key sequence length for attention computation. :type seqlen_k: Int32 :param window_size_left: Left-side sliding window size for attention masking. :type window_size_left: Optional[Int32] :param window_size_right: Right-side sliding window size for attention masking. :type window_size_right: Optional[Int32] :return: Tuple of (begin, end) tile idx for the trailing mask. :rtype: Tuple[Int32, Int32] rrF) rrrurzr}rrrr) r~r6rrMrrrr trip_startrrtrailing_mask_begintrailing_mask_ends rget_trailing_mask_idzFusedMask.get_trailing_mask_id s>((:X)XYYr_gjr_r--iJPXZbdtuu --         aL:a=069F rr rem_countc Bd}tj|tjkp|tjkrtj|duryt |||||||\} } tj|dur8t |||||||\} } | | kr| | z }n| | z dz}n| | z dz}ntj|tjkr||dzdkrd}nd}n|tj kr||z} |d|dz}t||ddz|dzdz }t |||||||}|dz| z|dz}|dz| z|dz}t||z |}||zS)a Calculate the number of masked trips for the trailing mask. This is used for blocks that need special handling due to masking. :param mask_type: Type of mask to use :type mask_type: utils.MaskType :param blk_coord: Block coordinates. :type blk_coord: cute.Coord :param tile_shape: Shape of the tile. :type tile_shape: cute.Shape :param seqlen_q: Query sequence length for attention computation. :type seqlen_q: Int32 :param seqlen_k: Key sequence length for attention computation. :type seqlen_k: Int32 :param window_size_left: Left-side sliding window size for attention masking. :type window_size_left: Optional[Int32] :param window_size_right: Right-side sliding window size for attention masking. :type window_size_right: Optional[Int32] :param rem_count: Remaining count from previous calculations. :type rem_count: Int32 :return: Number of masked trips. :rtype: Int32 rNrF) rrruryrzr}rrrxr{rrr)r~r6rrMrrrrrrrrrr block_startrrmask_start_trip mask_end_trips rget_masked_trailing_countz#FusedMask.get_masked_trailing_countjsH  i8+??n9PXPnCn o o5 F!"34"?@@ I9B9W9W$%::6#%6%&6d&BCCI;D;X;X!!"  ()<<8&(8+.>>>!25E!E!25H!H1!L.1DDqHF   X-C C D D F*Q-'1,, (9 9 9!)X!5 #A,A6KHy|a'7:a=&H1&LMMI"11 !J!,a4FF:VW=XO'!m0BBzRS}TM8*EEF !!rc t|||||||t|||||||z t|||||||dz }|S)a Calculate the number of unmasked trips for the current block. This represents the number of trips that don't require special masking treatment. :param mask_type: Type of mask to use :type mask_type: utils.MaskType :param blk_coord: Block coordinates. :type blk_coord: cute.Coord :param tile_shape: Shape of the tile. :type tile_shape: cute.Shape :param seqlen_q: Query sequence length for attention computation. :type seqlen_q: Int32 :param seqlen_k: Key sequence length for attention computation. :type seqlen_k: Int32 :param window_size_left: Left-side sliding window size for attention masking. :type window_size_left: Optional[Int32] :param window_size_right: Right-side sliding window size for attention masking. :type window_size_right: Optional[Int32] :return: Number of unmasked trips. :rtype: Int32 r)r}rrr)r~r6rrMrrrrs rget_unmasked_trip_countz!FusedMask.get_unmasked_trip_countsF  $ $ !  00 ! $11 !  %  < racc_qkindex_qkcNd}tj|tjkr||z }tjt j|D]}||\} } tj|dup|durltj|tjkrH||z} | dz| zdz | ks| |krtj ||<| |ks| |krtj ||<ntj|dur9| |z|z| krtj ||<| |ks| |krtj ||<ntj|dur9| |z|z | krtj ||<| |ks| |krtj ||<ndt| |z|z|} td| |z|z } | | ks| | krtj ||<| |ks| |krtj ||<tj|tj kr| |ks| |krtj ||<dS)a Apply the appropriate mask to the attention scores. This method modifies the attention scores (acc_qk) based on the mask type and the positions in the index tensor. :param mask_type: Type of mask to use :type mask_type: utils.MaskType :param acc_qk: Accumulated QK attention scores tensor. :type acc_qk: cute.Tensor :param index_qk: Index tensor containing position information. :type index_qk: cute.Tensor :param seqlen_k: Key sequence length for attention computation. :type seqlen_k: Int32 :param seqlen_q: Query sequence length for attention computation. :type seqlen_q: Optional[int] :param window_size_left: Left-side sliding window size for attention masking. :type window_size_left: Optional[int] :param window_size_right: Right-side sliding window size for attention masking. :type window_size_right: Optional[int] rNrF) rrrurzranger/r@r{rinfrrrx)r~rrrMrrrriindex_qindex_kr max_K_index min_K_indexs r apply_maskzFusedMask.apply_mask s~@  i8+II J J )(Fty0011 - -A'{ GW!"2$">"_BS[_B_`` 1%i83R&RSS1)1X)=&! (::Q>HHGW_L_L_%,[Lq (**g.A.A%,[Lq '(8D(@AA1'*;;gEE%,[Lq (**g.A.A%,[Lq '(9T(ABB 1'*::WDD%,[Lq (**g.A.A%,[Lq "%g&69J&JH"U"UK"%a6)9rsj """"""""''''''544444&&&&&&.d.d.d.d.d.d.d.dbtotototototototonP )PzP P PPPP*KKzK#KKKK$'# Z'#S#s]#'#'# (%S# *> >? '#'#'#'#^ ) ) ) ) )ty ) ) )^-^-^-^-^-^-^-^-^-^-r