)`ig+dZddlZddlmZddlmZddlmZmZm Z m Z ddl m Z ddl Z ddlmZdd lmZmZmZmZmZmZdd lmZmZdd lmZmZmZmZmZm Z m!Z!m"Z"dd l#m$Z$dd l#m%Z%ddl#m&Z&ddl#m'Z'ddl#m(Z(ddl#m)Z)ddl#m*Z*ddl#m+Z+ddl#m,Z,ddl#m-Z-ddl#m.Z.ddl/m0Z0ddl#m1Z1dZ2 ddl3Z3dZ2nM#e4$rYnFe5$r?Z6e7e68Z9e:ddDZ;e;sYdZ6[6ndZ6[6wwxYwddlZ>m?Z?m@Z@mAZAmBZBmCZCmDZDdZEd ZFd!e jGd"eHe7fd#ZIejJd$ZKe"d%d&gde jLddd'fd(e jMd)e jMd*e e jMd+e jNd,e e jMd-eOd.ed/fd0ZPe"d%d&gde jLddd'fd(e jMd)e jMd*e e jMd+e jNd,e e jMd-eOd.ed/fd1ZQddde jLd'fd(e jMd)e jMd,e e jMd-eOd*e e jMd+e jNd.ed/fd2ZRddde jLd'fd3ee7d(e jMd)e jMd,e e jMd-eOd*e e jMd+e jNd.ed/fd4ZSe!ePeQd5eReS6eddde jLd'fd(e jMd)e jMd,e e jMd-eOd*e e jMd+e jNd.ed/d7e jMfd8ZTe"d%d&gde jLd9fd:e jMd;e jMd*e e jMd+e jNd.ed9f d<ZUde jLd9fd:e jMd;e jMd*e e jMd+e jNd.ed9f d=ZVde jLd9fd3ee7d:e jMd;e jMd*e e jMd+e jNd.ed9f d>ZWe!d9eUieVeW6ede jLd9fd:e jMd;e jMd*e e jMd+e jNd.ed9d7e jMf d?ZXejJd@ZYejJdAZZejJdBZ[ejJdCZ\ejJdDZ]eedEdFeefedGdHdIfJZ^ejJdKZ_eedEdFeefedGdHdLfJZ`d(e jMd)e jMd,e jMd-eOd*e jMdMe jMdNee7d7dfdOZad(e jMd)e jMdPe jMdQe jMd*e jMdMe jMdNee7d7dfdRZbdSe7dTe7fdUZcejJdVZdejJdWZedXeffdYZgejJe jLdfdZe jNd[eOfd\Zhe dd(e jMd)e jMd,e jMd-eOd*e e jMd7e jMf d]ZiejJd^Zj dd_e jMd`e jMdae jMdbeOdcefdde jMdee e jMfdfZk dd_e jMd`e jMdae jMdbeOdcefdde jMdee e jMfdgZlGdhdiZmGdjdkeZndlZodmZpdnZqdardoe jsjtfdpZudZe jNfdqZvdZe jNfdrZwejJdsZxejJ ddueffdvZy ddueffdwZz ddueffdxZ{ejJdyZ|dzZ}dZe jNfd{Z~d|e jMd(e jMd)e jMd}e jMd~e jMd*e e jMde jNfdZdZdZdZede jLddfd(e jMd)e jMde e jMd+e jNd*e e jMd.edf dZde jLddddtfd(e jMd)e jMde jMde jMde e jMd+e jNd*e e jMdefdeOdueffdZde jLdddddtfd(e jMd)e jMde jMde jMde e jMd+e jNd*e e jMdefdeOdMe jMdueffdZdZde jLdddddfd(e jMd)e jMde jMde jMde e jMd+e jNd*e e jMdefdeOd.eddeOfdZe"gdde jLdddddfd(e jMd)e jMde jMde jMde e jMd+e jNd*e e jMdefdeOd.eddeOfdZe"d%d&gde jLdddddfd(e jMd)e jMde jMde jMde e jMd+e jNd*e e jMdefdeOd.eddeOfdZe"gdde jLdddddfd(e jMd)e jMde jMde jMde e jMd+e jNd*e e jMdefdeOd.eddeOfdZde jLdddddfd3ee7d(e jMd)e jMde jMde jMde e jMd+e jNd*e e jMdefdeOd.eddeOfdZdZeedEdEeefedddedddfJZeedEdEeefedddedddfJZe!eeedee6ede jLdddddfd(e jMd)e jMde jMde jMde e jMd+e jNd*e e jMdefdeOd.eddeOd7e jMfdZe"gd dd:e jMd;e jMde jMde jMdZe jNd*e e jMd.edfdZe"gd dd:e jMd;e jMde jMde jMdZe jNd*e e jMd.edfdZe"gd dd:e jMd;e jMde jMde jMdZe jNd*e e jMd.edfdZ dd:e jMd;e jMde jMde jMdZe jNd*e e jMd.edfdZ dd3ee7d:e jMd;e jMde jMde jMdZe jNd*e e jMd.edfdZe!eeedee6e dd:e jMd;e jMde jMde jMdZe jNd*e e jMd.edd7e jMfdZe"gd dd(e jMd)e jMd}e jMd~e jMde eddefde efefeffd*e e jMd+e e jNd.edfdZe"d%d&g dd(e jMd)e jMd}e jMd~e jMde eddefde efefeffd*e e jMd+e e jNd.edfdZ dd(e jMd)e jMd}e jMd~e jMde eddefde efefeffd*e e jMd+e e jNd.edfdZe!eedee dd(e jMd)e jMd}e jMd~e jMde eddefde efefeffd*e e jMd+e e jNd.edd7e jMfdZejJdZe"gd dd(e jMd)e jMd}e jMd~e jMde eddefd*e e jMd+e e jNfdZe!iee dd(e jMd)e jMd}e jMd~e jMde eddefd*e e jMd+e e jNd7e jMfdZe"gd dd(e jMd)e jMd}e jMd~e jMde jMde efefeffdeddefd*e e jMd+e e jNfdZe!iee dd(e jMd)e jMd}e jMd~e jMde jMde efefeffdeddefd*e e jMd+e e jNd7e jMfdZe"gd dd(e jMd)e jMd}e jMd~e jMde jMdefdefdefdefdeOd*e e jMd+e e jNfdƄZe!iee dd(e jMd)e jMd}e jMd~e jMde jMdefdefdefdefdeOd*e e jMd+e e jNd7e jMfdDŽZeZde jMfdȄZejJdɄZe"d%d&g dd(e jMd)e jMd}e jMd~e jMde jMde efefeffd*e e jMd+e e jNd7eOfd˄Ze!iee dd(e jMd)e jMd}e jMd~e jMde jMde efefeffd*e e jMd+e e jNfd̄Ze"d%d&g dd(e jMd)e jMd}e jMd~e jMde jMdefde efefeffd*e e jMd+e e jNd7eOfdτZe!iee dd(e jMd)e jMd}e jMd~e jMde jMdefde efefeffd*e e jMd+e e jNfdЄZejJdфZe dde jMde jMde e jMde e jMd*e e jMd+e e jNd7e jMfdքZdefdefdefdefd7e efefefff dڄZejJdۄZe jLdddtfd(e jMd)e jMd+e jNd*e e jMdefdueff d݄Ze jLdddtfd(e jMd)e jMde jMde jMd+e jNd*e e jMdMe jMdueffdބZd߄Zd(e jMd)e jMdPe jMdQe jMd*e jMdMe jMdNee7d7dfdZe"d%d&g dd:e jMd;e jMde jMde jMdZe jNd*e e jMd.edfdZdZe jNfdZ dd:e jMd;e jMde jMde jMdZe jNd*e e jMd.edfdZ dd3ee7d:e jMd;e jMde jMde jMdZe jNd*e e jMd.edfdZe!deiee6e dd:e jMd;e jMde jMde jMdZe jNd*e e jMd.edd7e jMfdZdS)a3 Copyright (c) 2024 by FlashInfer team. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. N)EnumSimpleNamespace)ListLiteralOptionalTuple)trtllm_low_latency_gemm)flashinfer_api) AutoTunerConstraintSpecDynamicTensorSpecOptimizationProfile TunableRunner TuningConfig)&get_last_power_of_2_num_tokens_bucketslast_positive_power_of_2)get_native_fp4_dtypeis_sm100a_supportedis_sm100f_supportedis_sm120a_supportedis_sm121a_supported LibraryErrorbackend_requirementsupported_compute_capability)gen_gemm_sm90_module)gen_gemm_module)gen_gemm_sm100_module)gen_gemm_sm120_module)!gen_gemm_sm120_module_cutlass_fp4)!gen_gemm_sm100_module_cutlass_fp4)!gen_gemm_sm100_module_cutlass_fp8)"gen_gemm_sm100_module_cutlass_bf16)gen_trtllm_gen_gemm_module)gen_tgv_gemm_sm10x_modulegen_deepgemm_sm100_module)get_cuda_version)#gen_fp8_blockscale_gemm_sm90_moduleFTc#(K|] }|tvVdSN) error_msg).0exts m/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/flashinfer/gemm/gemm_base.py r1Hs'EEc )EEEEEE)z.soz.dll)setup_cubin_loader)_get_cache_bufdetermine_gemm_backend get_indptr is_float8register_custom_opregister_fake_opget_compute_capabilityizecudnn FP4 GEMM with mxfp4 quantization is not supported on SM120 with cuDNN backend version < 9.14.0.device sm_versionc@t|\}}|dz|z}||vS)N )r:)r;r<majorminor device_archs r0_match_sm_versionrB^s/)&11LE5RZ%')K * $$r2ctfd}tdddtjdtjdtjdtjd tjd tjd tjd tjdtjd tjdt ddffd }t ddtjdtjdtjdtjd tjd tjd tjd tjdtjd tjdt ddfd}t||}|S)NcBGfddt}|S)Nc eZdZdeejdedeefdZ d deejdede dejffd Z d S) Lget_gemm_module..cublas_fp8_gemm_runner..CublasFp8GemmRunnerinputsprofilereturncdgSNrselfrGrHs r0get_valid_tacticsz^get_gemm_module..cublas_fp8_gemm_runner..CublasFp8GemmRunner.get_valid_tacticsks s r2Ftacticdo_preparationc tj}|\}}}} } } ||| || | || Sr,)torchcudacurrent_blas_handlebmm_fp8) rNrGrQrRkwargs cublas_handleabscale_ascale_boutworkspace_buffermodules r0forwardzTget_gemm_module..cublas_fp8_gemm_runner..CublasFp8GemmRunner.forwardssX!& > > @ @ @F=1gw-=q#w1A= r2NrPF __name__ __module__ __qualname__rrTTensorrintrOboolrar`sr0CublasFp8GemmRunnerrFjs U\* - c     !',  U\*  !%          r2rkr)rkr`s r0cublas_fp8_gemm_runnerz/get_gemm_module..cublas_fp8_gemm_runnerisH       -   .#"$$$r2z flashinfer::cutlass_segment_gemmy mutates_argsr_ all_problemsx_dataw_datay_datax_ldw_ldy_ld empty_x_dataweight_column_majorrIc D ||||||||| | dSr,)cutlass_segment_gemm) r_rqrrrsrtrurvrwrnrxryr`s r0r{z-get_gemm_module..cutlass_segment_gemmsE ##           r2c dSr,rL) r_rqrrrsrtrurvrwrnrxrys r0_fake_cutlass_segment_gemmz3get_gemm_module.._fake_cutlass_segment_gemms  r2)rmr{)rbuild_and_loadr8rTrgrir9r)rmr{r} _gemm_moduler`s @r0get_gemm_modulerds    - - / /F%%%%%8:#OOO , l       l  l l  < l "       PO 4899  ,  l         l  l  l   <  l  "       :9  #51L r2dgtgvrZr[r^ out_dtypebiaspdlbackend)cutlassrautoch|td|rtdt|dS)NzLYou cannot use the CUTLASS backend with a bias. Use the TGV backend instead.zFThe CUTLASS backend does not support PDL. Use the TGV backend instead.T) ValueError_validate_bf16_output_dtyperZr[r^rrrrs r0_cutlass_mm_bf16_requirementrsS  Z      T    *** 4r2cD|tjkrtddS)NzWYou cannot provide an output dtype to the TGV backend. Use the CUTLASS backend instead.T)rTbfloat16rrs r0_tgv_gemm_requirementrs-EN"" e    4r2c|jtjkrtd|jd|jtjkrtd|jd|-|jtjkrtd|jddS)N#First tensor has unsupported dtype . Only bfloat16 is supported.$Second tensor has unsupported dtype z"Bias tensor has unsupported dtype TdtyperTrr)rZr[rrr^rrs r0_check_mm_bf16_problem_sizers w%.   X!' X X X    w%.   Y17 Y Y Y    DJ%.88 Z Z Z Z    4r2suitable_backendscg}||rd|vr|dn2d|vr|dd|vr|d|S)Nrrappend) rrZr[rrr^rrheuristic_backendss r0_heuristic_func_mm_bf16rs| 3 % % %  % %e , , , ) ) )  % %i 0 0 0 % % %  % %e , , , r2)rr) common_checkheuristic_funcrIc |5tj|jd|jdf|j|}n|j|jd|jdfkr3t d|jd|jdfd|jd|j|jkr t d|jd|jd|j|krt d |d|jdt d t|j}|d kr tj }n?|d krtd g||dd |||}n!|dkrtdg|||||||}n|g}t||||||||S)a{MM BF16 Parameters ---------- a: torch.Tensor Input tensor, shape (m, k), bf16 in row-major layout. b: torch.Tensor Weight tensor, shape (k, n), bf16 in column-major layout. bias: Optional[torch.Tensor] Optional bias tensor, shape (n,). If provided, can only be used with the TGV backend. Defaults to ``None``. pdl: bool Whether to use persistant data loader mode. Can only be used with the TGV backend. Defaults to ``False``. out: Optional[torch.Tensor] Out tensor, shape (m, n), bf16 or fp16. If provided, can only be used with the CUTLASS backend. Defaults to ``None``. out_dtype: torch.dtype Output dtype, bf16 or fp16. If provided, can only be used with the CUTLASS backend. Defaults to ``torch.bfloat16``. backend: Literal["cutlass", "tgv", "auto"] The backend to use for the operation. Defaults to ``"tgv"``. ``"auto"`` allows selecting the best tactic from all available backends when autotune is enabled. Returns ------- torch.Tensor Out tensor, shape (m, n), bf16 or fp16 in row-major layout. Examples -------- >>> import torch >>> import flashinfer >>> # Using the TGV backend >>> a = torch.randn([48, 64], device="cuda", dtype=torch.bfloat16) >>> b = torch.randn([80, 64], device="cuda", dtype=torch.bfloat16).transpose(-2, -1) >>> bias = torch.randn([80], device="cuda", dtype=torch.bfloat16) >>> out = flashinfer.mm_bf16(a, b, bias=bias, pdl=True, backend="tgv") >>> out.shape torch.Size([48, 80]) >>> out.dtype torch.bfloat16 >>> # Using the CUTLASS backend >>> fp16_out = torch.empty([48, 80], device="cuda", dtype=torch.float16) >>> out = flashinfer.mm_bf16(a, b, out=fp16_out, out_dtype=torch.float16, backend="cutlass") >>> out.shape torch.Size([48, 80]) >>> out.dtype torch.float16 Nrr;r Output shape mismatch. Expected , got .!Output device mismatch. Expected Output dtype mismatch. Expected mm_bf16_workspacerrFr) rTemptyshaper;rrr4DEFAULT_WORKSPACE_SIZEmm_bf16suitable_auto_backendsrbf16_gemm_sm100) rZr[rrr^rrr_backendss r0rrsN {k WQZ $8    9QWQZ0 0 0_AGAJ 3K__SVS\___  : ! !QAHQQCJQQQ  9 ! !P9PPCIPPP &3QX&1 I  * KAtUCG   E  * GQ4c9g  9Aq$S*:HEEE Jr2rABc$t|dSNT)rrrr^rrs r0_cutlass_bmm_bf16_requirementrs *** 4r2c|jtjkrtd|jd|jtjkrtd|jddS)NrrrTrrs r0_check_bmm_bf16_problem_sizerso w%.   X!' X X X    w%.   Y17 Y Y Y    4r2c<g}d|vr|d|S)Nrr)rrrr^rrrs r0_heuristic_func_bmm_bf16rs0%%%!!),,, r2c |jd|jd|jdf}|tj||j|}n||j|krt d|d|jd|j|jkr t d |jd|jd|j|krt d |d|jdt d t|j}t||dd ||d g|S)aBMM BF16 Parameters ---------- A: torch.Tensor Input tensor, shape (b, m, k), bf16 in row-major layout. B: torch.Tensor Weight tensor, shape (b, k, n), bf16 in column-major layout. out: Optional[torch.Tensor] Out tensor, shape (b, m, n), bf16 or fp16, defaults to ``None``. out_dtype: torch.dtype Output dtype, bf16 (default) or fp16. backend: Literal["cutlass"] Backend to use, defaults to "cutlass". Returns ------- torch.Tensor Out tensor, shape (b, m, n), bf16 or fp16 in row-major layout. Examples -------- >>> import torch >>> import flashinfer >>> input = torch.randn([16, 48, 64], device="cuda", dtype=torch.bfloat16) >>> weight = torch.randn([16, 80, 64], device="cuda", dtype=torch.bfloat16).transpose(-2, -1) >>> out = flashinfer.bmm_bf16(input, weight) >>> out.shape torch.Size([16, 48, 80]) >>> out.dtype torch.bfloat16 rrr Nrrrrrrbmm_bf16_workspaceFr) rrTrr;rrr4rr)rrr^rrexpected_shaper_s r0bmm_bf16rs:hgaj!'!*agaj9N {k 8    9 & &U>UUUUU  : ! !QAHQQCJQQQ  9 ! !P9PPCIPPP &4ahAq$s,.cutlass_fp8_gemm_runner..CutlassFp8GemmRunnerrGrHrIcdgSNrPrLrMs r0rOzrget_gemm_sm120_module_cutlass_fp8..cutlass_fp8_gemm_runner..CutlassFp8GemmRunner.get_valid_tacticss t r2rPFrQrRc |\}}}}} } |dkr|dd} n|} |dkr*| jd} |jd} |jd}d}n4| jd} |jd} |jd}|jd}d}d}d}d}|| |}|||}|| k}||k}|s|s|}| }n;|dkr|}|r.cutlass_fp8_gemm_runner..CutlassFp8GemmRunner.forward.._pad_to_multiple?s\A-(:hFFr2rr;rMN)dim transposerrTnn functionalpad contiguouszerosrr;copy_rnumelviewexpandgemm_fp8_nt_groupwise)"rNrGrQrRrXrZr[r\r]r^r_ b_col_majorn_dimm_dimk_dim batch_size scale_gran_m scale_gran_n scale_gran_krn_paddedk_paddedneeds_n_paddingneeds_k_paddinga_paddedb_col_major_paddedb_underlying_padded out_padded scale_m_count scale_k_countscale_a_expanded scale_n_countscale_b_expandedr`s" r0razhget_gemm_sm120_module_cutlass_fp8..cutlass_fp8_gemm_runner..CutlassFp8GemmRunner.forwardsAG=1gw-=5577a<<#$++b""5"5KK#$K5577a<<'-a0EGAJEGAJE!"JJ'-a0EGAJEGAJE!"J " " GGG,+E<@@++E<@@"*e"3"*e"3&QQ H)4&&uuww!||#$*',x':'>'> ( 3 3 5 58e;K7L((H.3[%x0"-"3#.#5...* +6E66E6>:@@MMMM#$*',x':'>'> ( 3 3 5 58e;K7L((H/4k'8<"-"3#.#5///+ .A-J-J2r-R-R**111fuffuf+<=CCKPPP# %uuww!||%*["H-cj &&& &+['9#&:"%)&&& "%J==??a''"U*\9A=%%&M&. %.cutlass_fp8_gemm_runnersSZ Z Z Z Z Z Z =Z Z Z x$#%%%r2r)rrrr`s @r0!get_gemm_sm120_module_cutlass_fp8r sK# $ $F]&]&]&]&]&@  7   r2ct}|}t||Sr,)r%r~r3get_library_path)modops r0get_trtllm_gemm_modulers= $ & &C    Bs++--... Ir2cntfd}t|S)NcBGfddt}|S)Nc eZdZdeejdedeeffd Z d deejdede dejffd Z d S) `get_gemm_sm100_module_cutlass_fp8..cutlass_fp8_gemm_runner..CutlassFp8GemmRunnerrGrHrIc`ttSr,)listrangefp8_gemm_tactic_numrNrGrHr`s r0rOzrget_gemm_sm100_module_cutlass_fp8..cutlass_fp8_gemm_runner..CutlassFp8GemmRunner.get_valid_tactics' E&"<"<">">??@@@r2rPFrQrRc x|\}}}}} } ||dd||| | || SNrrP)fp8_gemmr) rNrGrQrRrXrZr[r\r]r^r_r`s r0razhget_gemm_sm100_module_cutlass_fp8..cutlass_fp8_gemm_runner..CutlassFp8GemmRunner.forwardsZAG=1gw-=KKB''$ r2Nrbrcrjsr0rrs AU\* A- Ac  A A A A A A!',   U\*  !%            r2rrlrs r0rzBget_gemm_sm100_module_cutlass_fp8..cutlass_fp8_gemm_runnersH       =   6$#%%%r2r)r#r~rrs @r0!get_gemm_sm100_module_cutlass_fp8r sO . 0 0 ? ? A AF&&&&&>  7   r2)rrrc|ddSNrrrLshapess r0r6!9R=r2dynamic_tensor_specsconstraint_specscntfd}t|S)NcBGfddt}|S)Nc eZdZdeejdedeeffd Z d deejdede dejffd Z d S) cget_gemm_sm100_module_cutlass_bf16..cutlass_bf16_gemm_runner..CutlassBf16GemmRunnerrGrHrIc`ttSr,)rrbf16_gemm_tactic_numrs r0rOzuget_gemm_sm100_module_cutlass_bf16..cutlass_bf16_gemm_runner..CutlassBf16GemmRunner.get_valid_tacticss' E&"="="?"?@@AAAr2rPFrQrRct|\}}}}}} ||dd|| ||Sr ) bf16_gemmr) rNrGrQrRrXrZr[_r^r_r`s r0razkget_gemm_sm100_module_cutlass_bf16..cutlass_bf16_gemm_runner..CutlassBf16GemmRunner.forwardsT5;11aC!1  KKB''$  r2Nrbrcrjsr0CutlassBf16GemmRunnerrs BU\* B- Bc  B B B B B B!',   U\*  !%            r2r"rl)r"r`s r0cutlass_bf16_gemm_runnerzDget_gemm_sm100_module_cutlass_bf16..cutlass_bf16_gemm_runnersH       M   2%$&&&r2)r#)r$r~r)r#r`s @r0"get_gemm_sm100_module_cutlass_bf16r$sO / 1 1 @ @ B BF'''''8 !9   r2c|ddSrrLrs r0rr*rr2r_ runner_namescg}t|j}d|vr3|td|vr:|t |j||s Jdtj } ||||||g} | d|t| \} } | | | dS)NrrNo suitable runners foundr rGrQ) rr;rr$r#get_tgv_gemm_sm10x_modulertgv_gemm_runnerr get choose_one_BF16_GEMM_SM100_TUNING_CONFIG) rZr[rrr^r_r&runners use_sm_100ftunerrGrunnerrQs r0rr0sG%ah//KL  9;;TTVVWWW  %ag{ ; ; K K M M    /////7 MOOED#s$4 5F%%& NFF F&((((((r2r\r]cFg}d|vr3|td|vr3|td|vr3|t d|vr!|t |s Jdtj}||||||g} | d|t| \} } | | | dS)N cutlass_sm10x cutlass_sm12xcublascudnnr(r r)) rr rrrrm_cudnn_gemm_fp8_runnerr r,r-_FP8_GEMM_SM100_TUNING_CONFIG rZr[r\r]r^r_r&r/r1rGr2rQs r0fp8_gemm_sm100r;Os0G,&&8::RRTTUUU,&&8::RRTTUUU<((??AABBB,-//000 /////7 MOOEGWc+; .cutlass_fp4_gemm_runner..CutlassFp4GemmRunnerc j|_dSr,)fp4_gemm_fp4_gemm_runner)rNr`s r0__init__zg_create_cutlass_fp4_gemm_module..cutlass_fp4_gemm_runner..CutlassFp4GemmRunner.__init__ts(.%%%r2rGrHrIc`ttSr,)rrfp4_gemm_tactic_numrs r0rOzp_create_cutlass_fp4_gemm_module..cutlass_fp4_gemm_runner..CutlassFp4GemmRunner.get_valid_tacticswr r2rPFrQrRc |\ }}}}} } } } } } |jtjkr4|jtjkr|tj}|jtjkr4|jtjkr|tj} ||j||j| | | || Sr,)rrTuint8 float8_e4m3fnrrCT)rNrGrQrRrXrZr[ a_descale b_descalealphar!r^r_r`s r0razf_create_cutlass_fp4_gemm_module..cutlass_fp4_gemm_runner..CutlassFp4GemmRunner.forward~s$ $7ek))ioAT.T.T )u{ ; ;I7ek))ioAT.T.T )u{ ; ;IqsIy{E3@PRX r2Nrb) rdrerfrErrTrgrrhrOrirarjsr0CutlassFp4GemmRunnerrAss 8 8 8 8 8 AU\* A- Ac  A A A A A A!',   U\*  !%         r2rOrl)rOr`s r0cutlass_fp4_gemm_runnerz@_create_cutlass_fp4_gemm_module..cutlass_fp4_gemm_runnerrsI% % % % % % % =% % % N$#%%%r2)rPr)r`r<r=rPs` r0_create_cutlass_fp4_gemm_modulerQos9(&(&(&(&(&T  7   r2cdt}t|ddS)z&Get the SM100/103/110 FP4 GEMM module.zflashinfer::cutlass_fp4_gemmcutlass_fp4_gemm)r"r~rQrjs r0!get_gemm_sm100_module_cutlass_fp4rTs6/ 0 0 ? ? A AF *.0B  r2cdt}t|ddS)z"Get the SM120/121 FP4 GEMM module.z"flashinfer::cutlass_fp4_gemm_sm120cutlass_fp4_gemm_sm120)r!r~rQrjs r0!get_gemm_sm120_module_cutlass_fp4rWs6/ 0 0 ? ? A AF *46N  r2sm_majorcr|dvrtS|dkrtStd|)N)r> zUnsupported SM major version: )rTrWr)rXs r0get_cutlass_fp4_gemm_moduler\sG80222 R0222D(DDEEEr2rr0crt||fd}t|S)a Get and build the TGV GEMM module for the specified dtype. Args: dtype: Data type for the GEMM operation (torch.bfloat16 or torch.float16) use_sm_100f: Whether to compile with SM100f flags (default: False), which makes the compiled kernel compatible with both B200 and B300 GPUs. However, it's only available with CUDA 12.9+. Returns: SimpleNamespace with the runner function cBGfddt}|S)Nc eZdZdeejdedeeffd Z d deejdede dejffd Z d S) Iget_tgv_gemm_sm10x_module..tgv_gemm_runner..TGVGemmRunnerrGrHrIcZj}tt|Sr,)tgv_gemm_tactic_numrr)rNrGrH tactic_fnr`s r0rOz[get_tgv_gemm_sm10x_module..tgv_gemm_runner..TGVGemmRunner.get_valid_tacticss)#6 E))++..///r2rPFrQrRc|^}}}}} } j} | ||||| || Sr,)tgv_gemmt) rNrGrQrRrXrZr[rrr^r!gemm_fnr`s r0razQget_tgv_gemm_sm10x_module..tgv_gemm_runner..TGVGemmRunner.forwardsO,2(1dCq !/qssuudFC=== r2Nrbrcrjsr0 TGVGemmRunnerr`s 0U\* 0- 0c  0 0 0 0 0 0!',   U\*  !%            r2rhrl)rhr`s r0r+z2get_tgv_gemm_sm10x_module..tgv_gemm_runnersE       M   6}r2)r+)r&r~r)rr0r+r`s @r0r*r*sU'uk : : I I K KF> '   r2ct|jddgstd|jtjtjfvrtd|jd|j|jkr td|jd|jd|:t j|jd |jd f|j|j }n|j|jd |jd fkr3td |jd |jd fd|jd|j|jkr td|jd|jd|j|jkr td|jd|jdg}t|j}| t|j| tj}d }tt!|fdt"t$ft'dddf} |||||g} |jtjkrdnd} || d|| | \} } | | | S)a+ Perform TGV GEMM on SM100 architecture with automatic dtype detection. Computes: A @ B + bias Args: a: First input tensor of shape (M, K) in row-major layout b: Second input tensor of shape (K, N) in column-major layout bias: Bias tensor of shape (N,) pdl: Whether to use PDL (persistent data loader), defaults to False out: Optional output tensor, shape (M, N), defaults to None. Returns: Output tensor of shape (M, N) in row-major layout Supported dtypes: - torch.bfloat16 - torch.float16 Note: - Requires SM100, SM103, or SM110 architecture - Input tensors a and b must have the same dtype - Tensor b is expected to be in column-major layout (transposed from typical PyTorch row-major) 100103z+TGV GEMM requires SM100, SM103 architecturezUnsupported dtype z*. Only bfloat16 and float16 are supported.z,Input tensors must have the same dtype. Got  and rNrrrrrrrrrrc|ddSrrLrs r0rz tgv_gemm_sm100..Isvay}r2rbf16fp16 _tgv_gemmr))rBr;rrrTrfloat16rrrrr*r+r r,rrrrrr-)rZr[rrr^r/r0r1a_tensor_index tuning_configrG dtype_strr2rQs r0tgv_gemm_sm100rusB QXu~ 6 6HFGGG wu~u}555 T T T T    w!' S17 S S S S S    {k WQZ $8'    9QWQZ0 0 0_AGAJ 3K__SVS\___  : ! !QAHQQCJQQQ  9  N17NN#)NNN G%ah//K NN,QWkBBRRTTUUU MOOEN !6(     ,,   M$D#s #F'U^33I%%  NFF 6 / / //r2c4ttdddtjdtjdtjdtjdtjd tjd tjd tjd tjd tjdtjdtjdt ddffd }t ddtjdtjdtjdtjdtjd tjd tjd tjd tjd tjdtjdtjdt ddfd}t|S)Nz%flashinfer::cutlass_segment_gemm_sm90)r_rnror_int_workspace_bufferrqrrrsrtx_stridew_stridey_stridernrx empty_y_dataryrIc H |||||||||| | | dSr,cutlass_segment_gemm_sm90)r_rwrqrrrsrtrxryrzrnrxr{ryr`s r0r~z7get_gemm_sm90_module..cutlass_segment_gemm_sm90`sK& ((            r2c dSr,rL) r_rwrqrrrsrtrxryrzrnrxr{rys r0_fake_cutlass_segment_gemm_sm90z=get_gemm_sm90_module.._fake_cutlass_segment_gemm_sm90s r2r})rr~r8rTrgrir9r)r~rr`s @r0get_gemm_sm90_modulerZs ! # # 2 2 4 4F/. , #l l       , , ,  < l l "        <=>> , #l l       , , ,  < l l "     ?> $ ";   r2rweightsrnw_column_majorr seg_indptrweight_indicesc|j}tj}tj} tj} || }||| }|r|dn|d} |r|dn|d} tj|df||} tj|| |}tj|| |}tj|| |}tj|| |}tj|| |}tj|| |}ddlm}||f| |||||||||||| | || ||||||fS)Nrr r)compute_sm80_group_gemm_args) r;rTint32int64tosizer triton.gemmr)rrrnrrrrr; prob_typeptr_typeld_typed_outd_inrqrrrsrt x_stride_data w_stride_data y_stride_datars r0#launch_compute_sm80_group_gemm_argsrsXF I{HkGx((J!'**844- BGLLOOO7<<??E, A7<<???',,q//D; AiOOOL [8F C C CF [8F C C CF [8F C C CFK '&IIIMK '&IIIMK '&IIIM::::::/ */    &  r2c|j}tj}tj} tj} || }||| }|r|dn|d} |r|dn|d} tj|df||} tj|| |}tj|| |}tj|| |}tj|| |}tj|| |}tj|| |}ddlm}||f| |||||||||||| | || ||||||fS)Nrr rr)compute_sm90_group_gemm_args) r;rTrrrrrrr)rrrnrrrrr;rr stride_typerrrqrrrsrtrrrrs r0#launch_compute_sm90_group_gemm_argsrsXF I{H+Kx((J!'**844- BGLLOOO7<<??E, A7<<???',,q//D; AiOOOL [8F C C CF [8F C C CF [8F C C CFK +fMMMMK +fMMMMK +fMMMM::::::/ */    &  r2ceZdZdZ ddejdeddfdZdejdejddfd Ze dd ejd ejd e d e de ejde ejde ejde ejdejfdZ e ZdS)SegmentGEMMWrapperaWrapper for segment GEMM kernels. Example ------- >>> import torch >>> from flashinfer import SegmentGEMMWrapper >>> # create a 1MB workspace buffer >>> workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda") >>> segment_gemm = SegmentGEMMWrapper(workspace_buffer) >>> seq_lens = torch.tensor([1, 2, 3, 4], dtype=torch.int64, device="cuda") >>> # create packed input tensor (10 = 1 + 2 + 3 + 4) >>> x = torch.randn(10, 128, device="cuda", dtype=torch.float16) >>> # create weight tensor with 4 weights, each with 128 input and 256 output channels, column major >>> weights = torch.randn(4, 256, 128, device="cuda", dtype=torch.float16) >>> # compute the segment GEMM >>> y = segment_gemm.run(x, weights, 4, True, seg_lens=seq_lens) >>> y.shape torch.Size([10, 256]) >>> y_ref_0 = torch.matmul(x[:1], weights[0].t()) >>> torch.allclose(y[:1], y_ref_0) True >>> y_ref_1 = torch.matmul(x[1:3], weights[1].t()) >>> torch.allclose(y[1:3], y_ref_1) True >>> y_ref_2 = torch.matmul(x[3:6], weights[2].t()) >>> torch.allclose(y[3:6], y_ref_2) True >>> y_ref_3 = torch.matmul(x[6:], weights[3].t()) >>> torch.allclose(y[6:], y_ref_3) True >>> >>> # another example with weight indices >>> weight_indices = torch.tensor([0, 1, 0, 1], dtype=torch.int64, device="cuda") >>> y = segment_gemm.run(x, weights, 4, True, seg_lens=seq_lens, weight_indices=weight_indices) >>> y.shape torch.Size([10, 256]) >>> y_ref_0 = torch.matmul(x[:1], weights[0].t()) >>> torch.allclose(y[:1], y_ref_0) True >>> y_ref_1 = torch.matmul(x[1:3], weights[1].t()) >>> torch.allclose(y[1:3], y_ref_1) True >>> y_ref_2 = torch.matmul(x[3:6], weights[0].t()) >>> torch.allclose(y[3:6], y_ref_2) True >>> y_ref_3 = torch.matmul(x[6:], weights[1].t()) >>> torch.allclose(y[6:], y_ref_3) True rfloat_workspace_bufferrrINcxtjdtj|j|_||_||_dS)aInitialize the wrapper. Parameters ---------- float_workspace_buffer : torch.Tensor The workspace buffer for the kernels, we use it for storing intermediate results in cutlass segment GEMM kernels. Encouraged size is 128MB. )irN)rTrint8r;_int_workspace_buffer_float_workspace_bufferr)rNrrs r0rEzSegmentGEMMWrapper.__init__Is?&+[ %*5K5R& & & "(>$ r2rwc"||_||_dS)aReset the workspace buffer. Parameters ---------- float_workspace_buffer : torch.Tensor The new float workspace buffer for the kernels. int_workspace_buffer : torch.Tensor The new int workspace buffer for the kernels. N)rr)rNrrws r0reset_workspace_bufferz)SegmentGEMMWrapper.reset_workspace_bufferZs(>$%9"""r2rrrryr^seg_lensrrc h||td|"t||}| tjdtj}|d} |r|dn|d} |Bt|r tj} n|j } tj | | f| |j }n)|j | | fkrtd| | fd |j tjd|j |j } tjd|j |j } |j d krt|j }n|j }|d krUt|||||||\}}}}}}}t!|j|j||||||||| | | nf|d krNt)|||||||\}}}}}}}t+|j||||||||| | ntd ||S)a Run the segment GEMM kernel. Compute the matrix multiplication between a batch of input tensor (with variable number of rows, but fixed number of columns) and a batch of weight tensor with fixed number of rows and columns: .. math:: y[i] = x[i] \times W[i] if :attr:`weight_indices` is provided, we will select the weight tensor based on the indices in the :attr:`weight_indices` tensor: .. math:: y[i] = x[i] \times W[\text{weight_indices}[i]] We use Ragged Tensor to represent the input tensor :attr:`x` and the output tensor :attr:`y`, and each x[i] is a segment of the concatenated tensor. Please see :ref:`Ragged Tensor tutorial ` for more details. We use a ``seg_len`` or ``seg_indptr`` tensor (either would work) to indicate the start and end of each segment, where the ``seg_indptr`` is the cumulative sum of the ``seg_lens`` tensor (with an additional 0 at the beginning): .. math:: \text{seg_indptr}[i] = \sum_{j=0}^{i-1} \text{seg_lens}[j], \quad \text{seg_indptr}[0] = 0 - If ``seg_lens`` is provided, then :attr:`x` has shape ``(sum(seg_lens), d_in)`` and :attr:`y` has shape ``(sum(seg_lens), d_out)``, where ``d_in`` is the number of columns of the input tensor and ``d_out`` is the number of columns of the output tensor. - If ``seg_indptr`` is provided, then :attr:`x` has shape ``(seg_indptr[-1], d_in)`` and :attr:`y` has shape ``(seg_indptr[-1], d_out)``. Parameters ---------- x : torch.Tensor The input tensor with shape ``(sum(seg_lens), d_in)``. weights : torch.Tensor The 3D weight tensor with shape ``(num_weights, d_in, d_out)`` if :attr:`weight_column_major` is ``False``, or ``(num_weights, d_out, d_in)`` if :attr:`weight_column_major` is ``True``. batch_size : int The number of segments. weight_column_major : bool Whether the weight tensor is column major. out : Optional[torch.Tensor] The output tensor, with shape ``(sum(seg_lens), d_out)``. If not provided, a new tensor will be created internally. seg_lens : Optional[torch.Tensor] The length of each segment, with shape ``(batch_size,)``, expects a 1D tensor of dtype ``torch.int64``. seg_indptr : Optional[torch.Tensor] The indptr of the segments, with shape ``(batch_size + 1,)``, expects a 1D tensor of dtype ``torch.int64``. If this is provided, then :attr:`seg_lens` will be ignored, otherwise ``seg_indptr`` will be computed internally from :attr:`seg_lens`. weight_indices : Optional[torch.Tensor] The indices of the weight tensor to be selected for each segment, with shape ``(batch_size,)``. Expects a 1D tensor of dtype ``torch.int64``. If this is provided, then the weight tensor will be selected based on the indices in this tensor. Returns ------- torch.Tensor The output tensor with shape ``(sum(seg_lens), d_out)``. Nz1Either seg_lens or seg_indptr should be provided.rrrr rz'Output tensor shape mismatch, expected rrsm90sm80zUnsupported gemm backend: )rr6rrTrrrr7rrrr;rrr5rrr~rrrrr{)rNrrrryr^rrrcumulative_batch_sizerrrxr{rrqrrrsrtrrr x_ld_data w_ld_data y_ld_datas r0runzSegmentGEMMWrapper.runisR   2PQQ Q  #HKKNN33J  !"[%+>>>N !q #6K QGLLOO ;|| $!N G +&.iCCy2E::: m>SUZ=[mmbebkmm{1AGAHEEE {1CIcjIII <6 ! !,QX66GGlG f  4#  ! " " < <,*#      4#     2 2*#    C'CCDD D r2)rNNNN)rdrerf__doc__rTrgstrrErr rhrirrrarLr2r0rrsF00fDJ&+l=@ " :&+l :JO, :  : : : :'++/-115hh <hh h " h el # h5<(hU\*h!.h hhh^hTGGGr2rc2eZdZdZdZdZdZdZdZdZ dZ d Z d S) UIDszUIDs for CUDNN graph tensorsrrr rrN) rdrerfrA_UIDB_UID ALPHA_UIDBLOCK_DESCALE_A_UIDBLOCK_DESCALE_B_UID A_SCALE_UID B_SCALE_UIDO_UIDrLr2r0rrs>&& E EIKK EEEr2rc2tstddS)z7Check if cuDNN is available and raise exception if not.zcuDNN is not available. Please install cuDNN to use FP8 GEMM functions. You can install it with: pip install nvidia-cudnn-cu12 nvidia-cudnn-frontendN)CUDNN_AVAILABLE RuntimeErrorrLr2r0_check_cudnn_availabilityr$s,   [     r2ct tj}tt|ddd\}}||fdkrt d|dn5#ttttf$r}t d|d}~wwxYw tj }|dkrt d |d dS#ttf$r}t d |d}~wwxYw) zCCheck if cuDNN FP4 support is available and raise exception if not.rNr )r z(cuDNN FP4 requires version 1.13+, found zH. Upgrade: pip install --upgrade nvidia-cudnn-cu12 nvidia-cudnn-frontendz= 91002, found z. Please upgrade cuDNN backend.zIUnable to determine cuDNN backend version. FP4 requires backend >= 91002.) rr7 __version__maprhsplitr ImportErrorAttributeErrorr IndexErrorbackend_version TypeError) version_strr?r@ers r0_check_cudnn_fp4_availabilityr-sS ' 3 1 1# 6 6rr :;; u 5>G # #Z;ZZZ  $ Z @ J    /11 U " "1o111  # " I & W   s0AA++BBB!,CC5 C00C5cHtj}d}d}||kp||kS)z;Check if cuBLAS backend for FP4 GEMM is available in cuDNN.icid)r7r)rCUDNN_VERSION_9_11_1CUDNN_VERSION_9_13_0s r0!_is_cublas_fp4_available_in_cudnnrNs:+--O  // 3 2 2r2streamct!ttjatjt|jtS)z(Create and return a cached cuDNN handle.) _cudnn_handlerr7 create_handle set_stream cuda_stream)rs r0_get_cudnn_handler_s@!###+--  ]F$6777 r2c`|tjtjfvrtd|ddS)6Validate that the output dtype is either bf16 or fp16.Unsupported output dtype: zN. Only torch.bfloat16 and torch.float16 are supported for FP8 GEMM operations.NrTrrqrrs r0_validate_fp8_output_dtyperisJ U^U]333 \ \ \ \   43r2c`|tjtjfvrtd|ddS)rrzO. Only torch.bfloat16 and torch.float16 are supported for BF16 GEMM operations.Nrrs r0rrrsJ U^U]333 ] ] ] ]   43r2cxtj| }tjt |5\}}| rtjjntjj}| d|||}| d|||}| d|||tj j }| d|||tj j }| ||d| gd }| tjj| ||| dgd }| tjj|||tjjd }| tjj|}| rt| d ddtjj}|d||tjj}|t$jj|dd | |t$jj|t$jj|t$jj|t$jj|t$jj|||tjj tjj!g| r$tEs|#dg|cdddS#1swxYwYdS)NrZnamerstride data_typer[block_descale_arrrrreordering_typeblock_descale_br dequant_a block_sizer dequant_bgemmcompute_data_typer global_scalerrr scale_mulrrZr[rc_finalTeng0)$rTrUcurrent_streamr7graphrrFP8_E4M3FP8_E8M0tensortensor_reorderingF8_128x4block_scale_dequantize set_data_typeFLOATmatmulmulset_uidrrvalueset_name set_outputrrrrrvalidatebuild_operation_graphcreate_execution_plans heur_moderrrdeselect_engines)a_shapea_strideb_shapeb_stridea_descale_shapea_descale_strideb_descale_shapeb_descale_strideab_typeo_typerr;alpha_is_not_none use_nvfp4rrr! scale_typea_cudnn_tensorb_cudnn_tensorblock_descale_a_cudnn_tensorblock_descale_b_cudnn_tensordequant_a_tensordequant_b_tensorc_tensorc_final_cudnn_tensorglobal_scale_cudnn_tensors r0%create_cudnn_execution_plans_fp4_gemmr!{s"Z & &v . .F &v.. / /P:E11:XU_--@X '(g&  '(g&  (-||"# !3< (4( ( $(-||"# !3< (4( ( $!77  (: 8   &&u'<=== 77  ("A 8   &&u'<===<<  #o3     u4555'  D(- # // )5)) % $)99 +"'/"7 $-$$ & - -dn.B C C C%%i00;;DAAOOPVWWWtz/000tz/000$,,T-E-KLLL$,,T-E-KLLL$$TZ%5666  ##%%% $$eo&79J%KLLL  -(I(K(K -  " "F8 , , ,aPPPPPPPPPPPPPPPPPPsM!N//N36N3rPrQct|||||||||| | | | | }||dkr||n||Sr)r! check_supportbuild_plan_at_index build_plans)r r r rrrrrrrrr;rNrrQrs r0 build_plans_cudnn_fp4_gemm_graphr&s& 2   E"  || !!&))))  Lr2c Htjj|t tjj|t tjj|tjj|tjj|i} |1|tj | tj j<| | kr8tj| |jtj}tj|j} |dkr'|| |t)| dS|| ||t)| dSNrrPhandle)rrrrrrrrrrTfloatrrget_workspace_sizerr;rIrUrexecuterexecute_plan_at_index) rrZr[rLrMrNrr_rQ variant_packrs r0execute_cudnn_gemm_fp4_graphr0 s_ !&&!5!7!788 !&&!5!7!788  &   &  ' L -2ZZ -D-D T^)*%":":"<"<<< ;  $ $ & &qxu{   Z & &qx 0 0F || l$4=Nv=V=V WWWWW ## *F;LV;T;T $     r2c Ftjj|tjj|tjj|tjj|tjj|i}|} || kr&tj | |j tj }tj |j } |dkr'|||t!| dS||||t!| dSr()rrrrrrrr,rrTrr;rIrUrr-rr.) rrZr[rLrMrr_rQr/workspace_sizers r0execute_cudnn_gemm_mxfp8_graphr31s ! !  &   &  ' L--//N.00 ; 185;   Z & &qx 0 0F || l$4=Nv=V=V WWWWW ## *F;LV;T;T $     r2c&ttj|}t jt |5\} } | d|||} | d|||} | dddtjj } | dddtjj }| d| | tjj }| d  tjj | d || tjj }| d ||tjj }| d d || t jj| t jj| t jj|t jj|t jj| | | tjjtjjg| | | cdddS#1swxYwYdS)aBuild a cuDNN graph for GEMM with per-tensor quantization. This function is cached to avoid rebuilding identical graphs. Args: a_shape: Shape of tensor A a_stride: Stride of tensor A b_shape: Shape of tensor B b_stride: Stride of tensor B a_type: Data type for input tensor A b_type: Data type for input tensor B o_type: Data type for output tensor Returns: cuDNN graph object rZrr[a_scalerb_scaler)rrrrc scale_mul_ar scale_mul_brTN)rrTrUrr7rrrrrrrrrrrrrrrrrrrrrr rFALLBACKr#r%)r r r ra_typeb_typerr;rrr!rra_scale_cudnn_tensorb_scale_cudnn_tensorc_cudnn_tensorc_after_scale_a_cudnn_tensorc_after_scale_b_cudnn_tensors r0(build_cudnn_gemm_with_per_tensor_q_graphrBTs( Z & &v . .F &v.. / /7:E1'(f&  '(f&   %||o+ ,   %||o+ ,   #o3 &   $$225?3HIII',yy"#o3 (1( ( $ (-yy*"#o3 (1( ( $ %--i88CCDIIWW     tz/000tz/000$$T%5%;<<<$$T%5%;<<<$,,TZ-=>>>  ##%%% $$eo&79Q%RSSS  o777777777777777777sJ*LL  L c tjj|tjj|tjj|tjj|tjj|i}tj |j }t|} | | kr8tj| |j tj}|||| dS)Nrr))rrrrrrrrTrUrr;rrr,rrIr-) rrZr[r5r6r workspacer/r cudnn_handles r0*execute_cudnn_gemm_with_per_tensor_q_graphrFs ! !   ' LZ & &qx 0 0F$V,,L5335555K  $ $ & &qxu{    MM, ,M?????r2c.|tjkrtjjS|tjkrtjjS|tjkrtjjS|tj krtjj Std|)NzUnsupported dtype: ) rTrr7rBFLOAT16rqHALFrJr float8_e5m2FP8_E5M2rrs r0#_torch_data_type_to_cudnn_data_typerLs{ '' %-  ## %% % %'' %# # #''6u66777r2rDr5r6torch_out_dtypec @tt|j||j|t |jt |jt ||j}t||||||||Sr,)rrBrrrLrr;rF)rDrZr[r5r6r^rMrs r0_cudnn_gemm_fp8rOs 4      +AG44+AG44+O<<   E/ q!WgsI Jr2c<Gddt}|S)Nc eZdZdeejdedeefdZ d deejdede dejfd Z d S) 2_cudnn_gemm_fp8_runner..CudnnFp8GemmRunnerrGrHrIcdgSrKrLrMs r0rOzD_cudnn_gemm_fp8_runner..CudnnFp8GemmRunner.get_valid_tacticss 3Jr2rPFrQrRc L|\}}}}} } t| ||||| | j| Sr,)rOr rNrGrQrRrXrZr[r\r]r^r_s r0raz:_cudnn_gemm_fp8_runner..CudnnFp8GemmRunner.forwards:=C 9Aq'7C)9 ,aGWc39 U U UJr2NrbrcrLr2r0CudnnFp8GemmRunnerrRs & ) #Y     #(  &  !  \      r2rVrl)rVs r0r8r8s=](    r2c|ddk}t|j}t|}t|dkr>|dd|d|||rdndxxdzcc<|rC|dxxdzcc<t t|dz D]}||xxdzcc<n2t t|dz D]}||xxdzcc<t|t|fS)Nrrr rrP)rrrleninsertrrtuple)packed_fp4_tensoris_column_major real_shape real_strideis r0%_get_real_fp4_shape_from_packed_uint8r`sp'..r22a7O'-..J(//1122K :!!Q1/5577888_,rr"---2--- B1s;''!+,,  A NNNa NNNN s;''!+,,  A NNNa NNNN *  u[11 22r2ct|j}t|}t|dkr|d||dd|ddk}|rdnd}|||zdksJ|||z||<||||z|d<n't|dkrnt d|dt |t |fS)Nr rrrrz&Unsupported block scale tensor shape: z, expected 2d or 3d.)rrrrXrYrrZ)block_scale_tensorrblock_scale_shapeblock_scale_strider\ expand_dims r0 _expand_block_scale_tensor_shaperfs@/56607799:: ""  J///!!!Q'''-33B771<)0QQq  ,z9Q>>>>(9*(E(S*% z *->z-J J 1   1 $ $  \5F \ \ \    # $ $e,>&?&? @@r2trtllm_low_latencyrNc tjf}d}|dkr|jd}|jd} ntd|d|d|9||vrtd |d|dtj|| f|j| }n|j|vrtd |jd|d|j|jd|jdfkr3td |jd|jdfd |jd |j|jkr td|jd |jd |&|j|krtd|d |jd |dkrt||||ntd|d|d|S)aVFP8 matrix multiplication. Parameters ---------- a: torch.Tensor Input tensor, shape (m, k), fp8 e4m3. b: torch.Tensor - When using "trtllm_low_latency" backend, Weight tensor, shape (k // block_size, n, block_size), fp8 e4m3 B needs to be pre-processed using `prepare_low_latency_gemm_weights`. block_size is 128 for e4m3. alpha: Optional[torch.Tensor] Scale tensor for the output, float. If None, defaults to 1.0 for no scaling. out_dtype: torch.dtype Output tensor data type. Default is torch.bfloat16. out: Optional[torch.Tensor] Output tensor, shape (m, n). If None, a new tensor will be allocated. backend: Literal["trtllm_low_latency"] Backend to use for computation. Default is "trtllm_low_latency". - "trtllm_low_latency": optimized for small M dimension. Returns ------- torch.Tensor Output tensor of shape (m, n) with dtype `out_dtype`. Examples -------- >>> import torch >>> from flashinfer import mm_fp8, prepare_low_latency_gemm_weights >>> m = 16 >>> n = 2560 >>> k = 32768 >>> a = torch.randn([m, k], device="cuda", dtype=torch.bfloat16) >>> a_fp8, a_inv_s = to_float8(a, dtype=torch.float8_e4m3fn) >>> b = torch.randn([n, k], device="cuda", dtype=torch.bfloat16) >>> b_fp8, b_inv_s = to_float8(b, dtype=torch.float8_e4m3fn) >>> prepared_b = prepare_low_latency_gemm_weights(b_fp8) >>> alpha = a_inv_s * b_inv_s >>> out = mm_fp8(a_fp8, prepared_b, alpha) >>> out.shape torch.Size([16, 2560]) )rgrgrrzUnsupported backend: z. Only z' are supported for FP8 GEMM operations.Nrrrrrrr)rTrrrrr;rr ) rZr[rNrr^rsupported_out_dtypessupported_backendsmns r0mm_fp8rm1sxt"N,0&&& GAJ GAJ PG P P& P P P    { 0 0 0VYVV,VVV k F8    90 0 0VSYVV,VVV  9QWQZ0 0 0]171:qwqz2J]]QTQZ]]]  : ! !QAHQQCJQQQ   SY)%;%;P9PPCIPPP &&&1eS1111 PG P P& P P P    Jr2rLrMrrc ,t|\} } t|\} } | d}t||\}}t||\}}t| | | | ||||tjjt |||j|du|| }|S)NrrQ)r`rfr&r7rFP4_E2M1rLr;)rZr[rLrMrNrr^rrrQ real_a_shape real_a_stride real_b_shape real_b_stridebatchexpanded_a_descale_shapeexpanded_a_descale_strideexpanded_b_descale_shapeexpanded_b_descale_striders r0_get_cudnn_fp4_gemm_graphr{s#H"J"JL-"G"J"JL- OE(E::87 )E::87 - ! !  +I66  T   E" Lr2c ht||||||||||  } t| ||||||| |  dS)N rZr[rLrMrNrr^rrrQrp)r{r0) rZr[rLrMrNrr^rrr_rQrs r0_cudnn_gemm_fp4r~sn &     E! q!Y 5#7GPVr2c<Gddt}|S)Nc eZdZdeejdedeefdZ d deejdede dejfd Z d S) 2_cudnn_gemm_fp4_runner..CudnnFp4GemmRunnerrGrHrIc |\ }}}}}}} } } } t||||||| | | d } | }tt|S)NrPr})r{get_execution_plan_countrr)rNrGrHrZr[rLrMrNrr^rrr_r num_planss r0rOzD_cudnn_gemm_fp4_runner..CudnnFp4GemmRunner.get_valid_tacticss"  .###%#   E6688Ii(()) )r2rPFrQrRc T|\ }}}}} } } } } }t||||| | | | | || dS)Nrp)r~)rNrGrQrRrXrZr[rLrMrNrr^rrr_s r0raz:_cudnn_gemm_fp4_runner..CudnnFp4GemmRunner.forward sl$         r2NrbrcrLr2r0CudnnFp4GemmRunnerrs" *&" *)" *#Y " *" *" *" *N#(   &  !   \       r2rrl)rs r0_cudnn_gemm_fp4_runnerrsFDDDDD]DDDL    r2ruse_8x4_sf_layout)r7trtllmrrc  |jdks |jdkrtd|jd|j|jd|jdkr+td|jdd|jd|jtjt hvs!|jtjt hvr td|jd|jd |jtjtjhvs|jtjtjhvr td |jd|jd |,|jtjkrtd |j|<| dkr$td | |tj tj fvrtd |d| r|dkrtd| s|dkrtddS)Nr zmm_fp4 accepts 2d tensors, got rlrrz1K dimension mismatch in mm_fp4. got a.shape[1] = z, b.shape[0] = z:a and b must have float4_e2m1fn_x2 packed into uint8. Got rzIa_descale and b_descale must have float8_e4m3fnx2 packed into uint8. Got z"alpha must be a float tensor, got zalpha must be a scalar, got rzN. Only torch.bfloat16 and torch.float16 are supported for FP4 GEMM operations.rnz$nvfp4 only supports block_size = 16. z$mxfp4 only supports block_size = 32.T) ndimrrrrTrIrrJr+rrrq rZr[rLrMrNrr^rrrrs r0_check_mm_fp4_problem_sizer9 sj v{{afkkR17RRRRSSSwqzQWQZ g g g[\[bcd[e g g    wu{$8$:$:;;;qw O@@ ,7 , ,!" , , ,      !4ek B B B <? < <)2 < < <    U[EK77KekKKLLL U[[]]a//G GGHHH777 \ \ \ \   AZ2%%?@@@ Ar))?@@@ 4r2)rrnxyc |rtd| sAt|jdgr+tjdkrt t tt|\} } t|\} }| d}t||\}}t||\}}t| | | |||||tj j t|||j|| }|dS)N6Only TRTLLM FP4 GEMM supports 8x4 scale factor layout.120ierT)rrBr;r7rr)CUDNN_FP4_MXFP4_SM120_CUDNN_VERSION_ERRORrr`rfr!rrqrLr#)rZr[rLrMrNrr^rrrrrrrsrtrurvrwrxryrzrs r0_cudnn_gemm_fp4_requirementrq s2SQRRR F ah 0 0F  ! # #e + +DEEE!### #H"J"JL-"G"J"JL- OE(E::87 )E::87 2 ! !  +I66    E   4r2c n| std|tjkrtd|ddS)N9Only cudnn and auto FP4 GEMM supports mxfp4 quantization.rzB. Only torch.bfloat16 is supported for TRTLLM FP4 GEMM operations.T)rrTrrs r0_trtllm_gemm_fp4_requirementr s\ VTUUUEN"" P P P P    4r2c J|rtd| stddS)NrrTrrs r0_cutlass_gemm_fp4_requirementr s9SQRRR VTUUU 4r2r7c tj} tr | dkrtjdkrd} nd} fd| DS)a Heuristic function for mm_fp4 backend selection. Routes to either cudnn or cutlass. Note: trtllm is not considered in the backend selection because it requires a specific input quantization (swizzling/shuffling) that differs from the preparation used for cudnn and cutlass backends. Logic for which comes first: - If cuda version is 12 - use cutlass. - If cuda version is 13 and cudnn version is less than 9.15 - use cutlass. - If cuda version is 13 and cudnn version is 9.15 or greater - use cudnn. rile)r7r)rr7cg|]}|v| SrLrL)r.r7rs r0 z*_heuristic_func_mm_fp4.. s$ D D D!Q2C-C-CA-C-C-Cr2)r)r?rr7r)rrZr[rLrMrNrr^rrrr cuda_majorcandidate_backendss` r0_heuristic_func_mm_fp4r si4"##)J2:++0E0G0G50P0P12 E D D D) D D DDr2c||zdz |z|zSrrL)rrns r0_pad_upr s UQY1  !!r2c:t|dddS)Nrrrs r0rr s76!9Q<33r2rc|ddSrKrLrs r0rr 6!9Q<r2c:t|dddS)Nrrrrs r0rr+ s76!9Q<55r2c|ddSrKrLrs r0rr0 rr2r7rrc |4tj|jd|jdf|j|}t dt |j} | dkr t j} n| g} t|j\} dfdfd d fd | D}tj }rtnt}||||||||| | g }| d |||\}}||| |S)a MM FP4 Parameters ---------- a: torch.Tensor Input tensor, shape (m, k), fp4 e2m1fn_x2 or uint8. b: torch.Tensor Mat2 tensor, shape (k, n), should be column major, fp4 e2m1fn_x2 or uint8. a_descale: torch.Tensor Block scale tensor for A, shape (m, k // block_size), float8_e4m3fn or uint8. b_descale: torch.Tensor Block scale tensor for B, shape (k, n // block_size), float8_e4m3fn or uint8. alpha: Optional[torch.Tensor] Global scale tensor, float scalar. out_dtype: torch.dtype Output dtype, bf16 or fp16. When ``backend="trtllm"``, only ``bf16`` is supported. out: Optional[torch.Tensor] Out tensor, shape (m, n), bf16 or fp16, defaults to ``None``. block_size: int Block size for FP4 quantization, only 16 and 32 are supported. 16 in case of nvfp4 quantization. 32 in case of mxfp4 quantization. use_8x4_sf_layout: bool Whether to use 8x4 scale factor layout or 128x4 scale factor layout, defaults to False. backend: Literal["cudnn", "trtllm", "cutlass", "auto"] Backend to use, defaults to ``"auto"``, which automatically selects the best backend between ``"cudnn"`` and ``"cutlass"`` based on the current CUDA and cuDNN versions. The ``"trtllm"`` backend is never selected when ``backend="auto"`` because it requires different weight preparation. use_nvfp4: bool Whether to use nvfp4 quantization or mxfp4 quantization, defaults to ``True``. See the ``block_size`` parameter for related constraints. Notes ----- When cudnn/cutlass backend is used, both a and b should quantized with nvfp4_quantize using the 128x4 scale factor layout and do_shuffle=False. When trtllm backend is used, b must be quantized with 128x4 layout and `do_shuffle=True`. a can be quantized with either 128x4 or 8x4 layout (controlled by `use_8x4_sf_layout`) and `do_shuffle=False`. Returns ------- out: torch.Tensor Out tensor, shape (m, n), bf16 or fp16. Examples -------- >>> import torch >>> from flashinfer import nvfp4_quantize, mm_fp4, SfLayout >>> a = torch.randn([48, 128], device="cuda", dtype=torch.bfloat16) >>> b = torch.randn([256, 128], device="cuda", dtype=torch.bfloat16) >>> a_global_sf = (448 * 6) / a.float().abs().nan_to_num().max() >>> b_global_sf = (448 * 6) / b.float().abs().nan_to_num().max() >>> a_fp4, a_sf = nvfp4_quantize(a, a_global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False) >>> b_fp4, b_sf = nvfp4_quantize(b, b_global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=True) >>> out = mm_fp4(a_fp4, b_fp4.T, a_sf, b_sf.T, 1.0/(a_global_sf * b_global_sf), torch.bfloat16, None, backend="trtllm") >>> out.shape torch.Size([48, 256]) Nrrrmm_fp4_workspacerctSr,)rrLr2r0rzmm_fp4.. s /11r2cFtSr,)get_trtllm_fp4_gemm_moduletrtllm_fp4_gemm_runner)rsr0rzmm_fp4.. s!466MM   r2cFtSr,)r\rP)r?sr0rzmm_fp4.. s6u==UUWWr2rc0g|]}|SrLrL)r. cur_backendbackend_to_runner_factorys r0rzmm_fp4.. s(TTTK5(577TTTr2rCr))rTrrr;r4rmm_fp4rr:r r,_MM_FP4_TUNING_CONFIG_8x4_MM_FP4_TUNING_CONFIG_128x4r-)rZr[rLrMrNrr^rrrrr_rr!r/r1rsrGr2rQrr?s ` @@r0rr6 sot {k WQZ $8    &2AH &09&ah//HE121    XWWW !!UTTT8TTTG MOOE&7W!!>> import torch >>> import torch.nn.functional as F >>> import flashinfer >>> def to_float8(x, dtype=torch.float8_e4m3fn): ... finfo = torch.finfo(dtype) ... min_val, max_val = x.aminmax() ... amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12) ... scale = finfo.max / amax ... x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max) ... return x_scl_sat.to(dtype), scale.float().reciprocal() >>> >>> input = torch.randn([16, 48, 64], device="cuda", dtype=torch.bfloat16) >>> input_fp8, input_inv_s = to_float8(input, dtype=torch.float8_e4m3fn) >>> # column major weight >>> weight = torch.randn([16, 80, 64], device="cuda", dtype=torch.bfloat16).transpose(-2, -1) >>> weight_fp8, weight_inv_s = to_float8(weight, dtype=torch.float8_e4m3fn) >>> out = flashinfer.bmm_fp8(input_fp8, weight_fp8, input_inv_s, weight_inv_s, torch.bfloat16) >>> out.shape torch.Size([16, 48, 80]) >>> out.dtype torch.bfloat16 Nrrr rbmm_fp8_workspacerrr7) rTrrr;r4rrWrrrr;) rrrrrr^rr_rs r0rWrW! sV {k WQZQWQZ 08    &3QX&1 I  * KAwW   G   991a'30@(KKK Jr2)rrrrrrrrscale_major_moderKmma_smscale_granularity_mnk)rrc (|tddS)Nz'scale_major_mode is required in CUTLASSTr rZr[r5r6rrrr^rrs r0*_cutlass_gemm_fp8_nt_groupwise_requirementr sBCCC 4r2c p|dkrtd|jddkrtddS)Nrz5scale_granularity_mnk must be (1, 128, 128) in TRTLLMrz#a.shape[1] must be >= 256 in TRTLLMT)rrrs r0)_trtllm_gemm_fp8_nt_groupwise_requirementr sC --PQQQwqzC>??? 4r2c L|jdks |jdkrtd|jd|j|jd|jdkr+td|jdd|jd||p tj}n|j}t |dS)Nr zShape mismatch. a.shape = z , b.shape = rShape mismatch. a.shape[1] = z, b.shape[1] = T)rrrrTrrrrs r0)_check_gemm_fp8_nt_groupwise_problem_sizer s v{{afkkTagTT17TTUUUwqzQWQZ SAGAJ S Sqwqz S S    {/ I y))) 4r2)rc rtdt|j} ||p tj}n|j}|3tj|jd|jd|j|}| dkrt|jst|jr"tj | |||||g||Rn~t|jr$tj | |||||g|||RnFtd|j| dkr)t| ||||d|dd |S) a Performs matrix multiplication with FP8 data types using groupwise scaling. This function implements a GEMM operation that allows for fine-grained control over scale granularity across different dimensions. Currently only supported on NVIDIA Blackwell architecture. Parameters ---------- a: torch.Tensor Row-major input tensor shape (m, k), fp8 e4m3 or fp8 e5m2. b: torch.Tensor Column-major input tensor shape (n, k), fp8 e4m3 or fp8 e5m2. a_scale: torch.Tensor if the backend is ``cutlass``: Column-major scale tensor for a, shape ``(m, k // block_size)`` if scale_major_mode is ``K`` or shape ``(k // block_size, m)`` if scale_major_mode is ``MN`` if the backend is ``trtllm``: scale_major_mode should be None, the scale tensor should be (m, k // block_size), contiguous on the first dimension b_scale: torch.Tensor if the backend is ``cutlass``: Row-major scale tensor for b, shape ``(n // block_size, k // block_size)`` if scale_major_k is ``K`` or shape ``(k // block_size, n // block_size)`` if scale_major_mode is ``MN`` if the backend is ``trtllm``: scale_major_mode should be None, the scale tensor should be (k // block_size, n // block_size), contiguous on the first dimension scale_granularity_mnk: Tuple[int, int, int] The granularity of the scale tensor, (m_granularity, n_granularity, k_granularity). scale_major_mode: Literal["MN", "K"] The layout mode of scale tensor, `MN` for MN-major scale with shape of ``(k // block_size, *)`` and `K` for K-major scale with shape of ``(*, k // block_size)`` mma_sm: int How many SMs to use for the MMA operation, must be 1 or 2. 2 is faster when number of rows (M) per group is large (>= 256). out: Optional[torch.Tensor] Output tensor, shape (m, n). If not specified, we will create an output tensor explicitly. out_dtype: Optional[torch.dtype] If out is not specified, we will create an output tensor with this dtype. Defaults to ``torch.bfloat16``. backend: Literal["cutlass", "trtllm"] The backend to use for the operation. Defaults to ``"cutlass"``. Returns ------- out: torch.Tensor Output tensor, shape (m, n). Notes ----- The ``m`` should be padded to a multiple of 4 before calling this function, to accommodate the kernel's requirement. gemm_fp8_nt_groupwise_workspaceNrrrz!Unsupported device for FP8 GEMM: rFrP)r4rr;rTrrrrrrrrrrrr trtllm_gemm) rZr[r5r6rrrr^rrr_s r0rr sd&)+A18 {/ I  {k GAJ GAJ8    ) qx ( ( M,?,I,I M 9 ! # # 9  ' !     ! * * M 9 ! # # 9  ' !      KKKLL L H    ,,         Jr2ct}|t|ddtffd }t |S)NTrcDGfddt}||S)NceZdZddeffd Zdeejdedee ffd Z ddeejd e d efd Z d S)Wget_trtllm_fp4_gemm_module..trtllm_fp4_gemm_runner..TrtllmFp4GemmRunnerTrc.j|_||_dSr,)rrD_use_8x4_sf_layout)rNrrs r0rEz`get_trtllm_fp4_gemm_module..trtllm_fp4_gemm_runner..TrtllmFp4GemmRunner.__init__g s(*%*;'''r2rGrHrIc $d}d}||}||}|d}|d}|ddz} |\ }}} } } } }} } }d}d}t||| |||jS)Nrr r)get_opt_shapesrtrtllm_gemm_tacticsr)rNrGrHrrb_tensor_indexrZr[rkrlkrLrMrNr!r^r_ type_e2m1 type_bf16rs r0rOziget_trtllm_fp4_gemm_module..trtllm_fp4_gemm_runner..TrtllmFp4GemmRunner.get_valid_tacticsk s "#!"**,,^<**,,^<aDaDaD1H $  **1aIt7Nr2rPFrQrRc x|\ }}}}} } } } } } || ||j||j| | |j| | Sr,)rDrKr) rNrGrQrRrXrZr[rLrMrNr!r^r_s r0raz_get_trtllm_fp4_gemm_module..trtllm_fp4_gemm_runner..TrtllmFp4GemmRunner.forward sk$ $%%$CK+    r2NTrb) rdrerfrirErrTrgrrhrOra)rsr0TrtllmFp4GemmRunnerrf s < <$ < < < < < < U\* - c       H!',   U\*  !%       r2rrl)rrrs r0rz:get_trtllm_fp4_gemm_module..trtllm_fp4_gemm_runnere sVD D D D D D D -D D D L#"#4555r2)rr)r%r~r3rrir)rrrs @r0rr_ s $ & &C    Bs++--...G6G6$G6G6G6G6G6G6T 5   r2rc jt||||||d||d t||||||d||d dS)Nrrrr)rr^rrT)rrrZr[r5r6rrr^rs r0+_check_gemm_fp8_nt_blockscaled_problem_sizer sw.  -     /  -      4r2c 2t||||d|||| S)zPerforms matrix multiplication with FP8 data types using block-scaled scaling. Block-scaled scaling is a special case of groupwise scaling where the scale granularity is (128, 128, 128). r)rrrr^r)rrs r0gemm_fp8_nt_blockscaledr s6( !  -)     r2m_indptrc |jtjtjfvrt d|j|jtjtjfvrt d|j|jtjfvrt d|j|jtjfvrt d|j|jtjfvrt d|j|dvrt d||dvrt d ||jd } |jd } || tj} nm| |j} |j|jd | fkr't d |jd|jd | f|j| krt d|jd| t| |jd | kr t d|jd d| | dzd krt d| | dzd krt d| |jd d z } t|j st|j r| d krtddS)Nz#a must be a float8 tensor, but got z#b must be a float8 tensor, but got z*a_scale must be a float32 tensor, but got z*b_scale must be a float32 tensor, but got )m_indptr must be a int32 tensor, but got rz5scale_major_mode must be either 'MN' or 'K', but got rr &mma_sm must be either 1 or 2, but got rr rzShape mismatch. out.shape = z, (a.shape[0], n) = zdtype mismatch. out.dtype = z, out_dtype = rz, k = rz#n must be a multiple of 8, but got rnz$k must be a multiple of 16, but got zRgroup_gemm_fp8_nt_groupwise has correctness issues for num_groups > 1 on SM120/121T)rrTrJrJrfloat32rrrrrr;rr) rZr[r5r6rrrrr^rrlr num_groupss r0/_check_group_gemm_fp8_nt_groupwise_problem_sizer s wu*E,=>>>HqwHHIIIwu*E,=>>>HqwHHIII}U]O++UgmUUVVV}U]O++UgmUUVVV~ek]**UX^UUVVV{** VDT V V   VJ&JJKKK  A  A {  I   I 9Q ' '_sy__qwWXz[\o__  9 ! !SsySS SS y)))wqzQNNN1NNOOO1uzzBqBBCCC2v{{CCCDDD"Q&J18$$(;AH(E(E >>d  4r2c Ftdt|j} tdt|j} || tj} n | |j} |jd} |jd} |jd| f}|tj|| |j}t|jst|jr&tj | | ||||||| | g ||Rn;t|jr'tj | | ||||||| | g |||R|S)aPerform group GEMM with FP8 data types using groupwise scaling. Currently only supported on NVIDIA Blackwell architecture. Parameters ---------- a: torch.Tensor Row-major input tensor shape ``(cum_m, k)``, data type is ``torch.float8_e4m3fn`` or ``torch.float8_e5m2``. ``cum_m`` is the cumulative sum of the segment lengths. b: torch.Tensor Column-major input tensor shape ``(batch_size, n, k)``, data type is ``torch.float8_e4m3fn`` or ``torch.float8_e5m2``. a_scale: torch.Tensor Column-major scale tensor for a, shape ``(cum_m, k // block_size)`` if scale_major_mode is ``K`` or shape ``(k // block_size, cum_m)`` if scale_major_mode is ``MN``, data type is ``torch.float32``. b_scale: torch.Tensor Row-major scale tensor for b, shape ``(batch_size, n // block_size, k // block_size)`` if scale_major_mode is ``K`` shape ``(batch_size, k // block_size, n // block_size)`` if scale_major_mode is ``MN``, data type is ``torch.float32``. m_indptr: torch.Tensor The indptr of the segment lengths, shape ``(batch_size + 1,)``, data type is ``torch.int32``. Element element in ``m_indptr`` must be a multiple of 4. scale_granularity_mnk: Tuple[int, int, int] The granularity of the scale tensor, (m_granularity, n_granularity, k_granularity). scale_major_mode: Literal["MN", "K"] The layout mode of scale tensor, `MN` for MN-major scale with shape of ``(k // block_size, *)`` and `K` for K-major scale with shape of ``(*, k // block_size)`` mma_sm: int How many SMs to use for the MMA operation, must be 1 or 2. 2 is faster when number of rows (M) per group is large (>= 256). out: Optional[torch.Tensor] The output tensor, shape ``(cum_m, n)``. If not specified, we will create an output tensor explicitly. out_dtype: Optional[torch.dtype] The data type of the output tensor, must be ``torch.bfloat16`` or ``torch.float16``. Returns ------- out: torch.Tensor The output tensor, shape ``(cum_m, n)``. Notes ----- Each value in ``m_indptr`` should be padded to a multiple of 4 before calling this function, to accommodate the kernel's requirement. )group_gemm_fp8_nt_groupwise_int_workspace+group_gemm_fp8_nt_groupwise_float_workspaceNrr rr)r4rr;rTrrrrrrrgroup_gemm_fp8_nt_groupwiserr)rZr[r5r6rrrrr^rrwrrlr out_shapes r0rrB sJ*35KQX,57Mqx {  I   I  A  AQI {k)9QXFFF18$$ (;AH(E(E ;; "      #    QX & & ;; "      #         Jr2rtile_mtile_ntile_kswap_abc |jtjtjfvrt d|j|jtjkrt d|j|jtjkrt d|j|jtjkrt d|j|jtjkrt d|j|dvrt d||dvrt d ||d vrt d ||d vrt d || dvrt d| | | tj} n | | j} | tjtjfvrt d| |j ddz } |j d| kr t d|j dd| |j d} |j ddz}|j d|kr t d|j dd|d}d}| |zdkrt d|d| ||zdkrt d|d||j d| f}| J| j |krt d|d| j | j| krt d | d| jd!S)"Nz9a must be a float8_e4m3fn or float8_e5m2 tensor, but got z"b must be a uint8 tensor, but got z(a_scale must be a uint8 tensor, but got z(b_scale must be a uint8 tensor, but got rrr)rztile_m must be 128, but got )@rrz3tile_n must be one of [64, 128, 192, 256], but got )rrz*tile_k must be either 128 or 256, but got )TFz)swap_ab must be a boolean value, but got zBout_dtype must be either torch.bfloat16 or torch.float16, but got rrzMb.shape[0] must equal num_groups (m_indptr.shape[0] - 1), but got b.shape[0]=z , num_groups=r z,a.shape[1] must equal k, but got a.shape[1]=z, k=rrzn must be a multiple of z , but got n=zk must be a multiple of z , but got k=zout.shape must be z , but got zout.dtype must be T) rrTrJrJrrIrrrqr)rZr[r5r6rrrrr r r^rrrlralign_nalign_krs r07_check_group_gemm_mxfp8_mxfp4_nt_groupwise_problem_sizer st wu*E,=>>> Q Q Q    w%+GagGGHHH} ##SGMSSTTT} ##SGMSSTTT~$$UX^UUVVV VJ&JJKKK U@@@AAA (((WvWWXXX ZNfNNOOOm##NWNNOOO {  I   I777 \QZ \ \   "Q&JwqzZ B\]\cde\f B Bu B B     A  QA wqzQ N171: N N1 N N   GG7{aLGLLLLMMM7{aLGLLLLMMMQI  9 ! !R)RRsyRRSS S 9 ! !R)RRsyRRSS S 4r2c tdt|j} tdt|j} | | tj} n | | j} |jd}|jddz}|jd|f}| tj|| |j} t | | ||||| |||||||| | S)aPerform group GEMM with MXFP4 data types using groupwise scaling. Currently only supported on NVIDIA Blackwell architecture. Parameters ---------- a: torch.Tensor Row-major input tensor, shape ``(cum_m, k)``, data type is ``torch.float8_e4m3fn`` or ``torch.float8_e5m2``. ``cum_m`` is the cumulative sum of the segment lengths. b: torch.Tensor Column-major input tensor, shape ``(batch_size, n, k // 2)``, data type is ``torch.uint8``. a_scale: torch.Tensor Column-major scale tensor for a, shape ``(cum_m_padded, k // 32)``, data type is ``torch.uint8``. b_scale: torch.Tensor Row-major scale tensor for b, shape ``(batch_size, n_padded, k // 32)``, data type is ``torch.uint8``. m_indptr: torch.Tensor The indptr of the segment lengths, shape ``(batch_size + 1,)``, data type is ``torch.int32``. Element element in ``m_indptr`` must be a multiple of 4. mma_sm: int How many SMs to use for the MMA operation, must be 1 or 2. 2 is faster when number of rows (M) per group is large (>= 256). tile_m: int The tile size for the M dimension, must be 128. tile_n: int The tile size for the N dimension, must be 64, 128, 192, or 256. tile_k: int The tile size for the K dimension, must be 128 or 256. swap_ab: bool Whether to swap the A and B tensors. out: Optional[torch.Tensor] The output tensor, shape ``(cum_m, n)``. If not specified, we will create an output tensor explicitly. out_dtype: Optional[torch.dtype] The data type of the output tensor, must be ``torch.bfloat16`` or ``torch.float16``. Returns ------- out: torch.Tensor The output tensor, shape ``(cum_m, n)``. Notes ----- Each value in ``m_indptr`` should be padded to a multiple of 4 before calling this function, to accommodate the kernel's requirement. +group_gemm_mxfp4_nt_groupwise_int_workspace-group_gemm_mxfp4_nt_groupwise_float_workspaceNrr rr) r4rr;rTrrrrrgroup_gemm_mxfp4_nt_groupwise)rZr[r5r6rrrrr r r^rrwrrlrrs r0#group_gemm_mxfp8_mxfp4_nt_groupwisersR*57Mqx,7   {  I   I  A  QAQI {k)9QXFFF99     " Jr2cddlm}|jddz }|dd|ddz }|dz|dzdzz }tjtjd|j|j |f}|d|j }tj|df|j|j }tj|df|j|j }||f||||||fS) Nr )compute_padding_mappingrrrPrr)rr)rrr) rrrrTcatrr;rcumsum)rrrrkpadded_m_indptrm_rank padded_m_ranks r0pad_indptr_to_multiple_of_4rs655555"Q&J x}$A AQ! AiT!(!'!R!R!RTU VWWO%,,/:O,PPO [(2,x W W WFK "x~hoM+ZM*/6= M ))r2c"t}|Sr,r'rjs r0get_deepgemm_sm100_modulers & ( (F Mr2 m_indicescddlm}|A|p tj}tj|jd|jd||j}|||f||f|||S)Nr)4_check_group_deepgemm_fp8_nt_contiguous_problem_sizerr)flashinfer.deep_gemmr"rTrrrr;) rZr[r5r6r rr^rr"s r03_check_group_deepgemm_fp8_nt_groupwise_problem_sizer$s {/ k!'!*agaj !(SSS ? ? G q'lC4I  r2cddlm}|A|p tj}tj|jd|jd||j}|||f||f||||S)aPerform grouped matrix multiplication with FP8 data types using DeepGEMM backend. This function performs a grouped GEMM operation where each group in tensor `b` is multiplied with the corresponding rows in tensor `a`. The grouping is determined by the `m_indices` tensor, which specifies which group each row belongs to. This is particularly useful for scenarios like mixture of experts (MoE) where different tokens are routed to different experts. The operation can be conceptualized as: >>> for i in range(num_groups): >>> row_slice = slice(i * m_per_group, (i + 1) * m_per_group) >>> output[row_slice] = a[row_slice] @ b[i].T Currently only supported on NVIDIA Blackwell (SM100) architecture. Parameters ---------- a : torch.Tensor Input tensor A of shape ``(m, k)`` with FP8 data type (``torch.float8_e4m3fn``). This tensor contains all rows that will be multiplied with different groups in `b`. b : torch.Tensor Input tensor B of shape ``(batch_size, n, k)`` with FP8 data type (``torch.float8_e4m3fn``). Each slice ``b[i]`` represents a different group/expert that will be multiplied with the corresponding rows in `a`. a_scale : torch.Tensor Scaling factors for tensor `a` of shape ``(m, k // block_size)`` with ``torch.float32`` dtype. These are typically generated from per-token quantization of the original float32 tensor. b_scale : torch.Tensor Scaling factors for tensor `b` of shape ``(batch_size, n // block_size, k // block_size)`` with ``torch.float32`` dtype. These are typically generated from per-block quantization of the original float32 tensor for each group. m_indices : torch.Tensor Group assignment tensor of shape ``(m,)`` with ``torch.int32`` dtype. Each element specifies which group (index into `b`) the corresponding row in `a` belongs to. For example, if ``m_indices[i] = j``, then row ``i`` in `a` will be multiplied with group ``j`` in `b`. scale_granularity_mnk : Tuple[int, int, int], optional The granularity of the scaling factors as ``(m_granularity, n_granularity, k_granularity)``. Default is ``(1, 128, 128)`` which means per-token scaling for `a` and 128x128 block scaling for `b`. out : Optional[torch.Tensor], optional Pre-allocated output tensor of shape ``(m, n)``. If not provided, a new tensor will be created. out_dtype : Optional[torch.dtype], optional Data type of the output tensor. If `out` is provided, this parameter is ignored. Default is ``torch.bfloat16``. Returns ------- torch.Tensor Output tensor of shape ``(m, n)`` containing the results of the grouped matrix multiplication. Examples -------- >>> import torch >>> from flashinfer.gemm import group_deepgemm_fp8_nt_groupwise >>> from flashinfer.utils import per_token_cast_to_fp8, per_block_cast_to_fp8 >>> >>> # Setup: 2 groups, 128 tokens per group, 4096 hidden size, 2048 expert size >>> m_per_group, n, k = 128, 2048, 4096 >>> group_size = 2 >>> m = m_per_group * group_size >>> >>> # Create float32 inputs >>> a_f32 = torch.randn(m, k, device="cuda", dtype=torch.float32) >>> b_f32 = torch.randn(group_size, n, k, device="cuda", dtype=torch.float32) >>> >>> # Quantize to FP8 with appropriate scaling >>> a_fp8, a_scale = per_token_cast_to_fp8(a_f32) >>> b_fp8 = torch.empty_like(b_f32, dtype=torch.float8_e4m3fn) >>> b_scale = torch.empty((group_size, n // 128, k // 128), device="cuda", dtype=torch.float32) >>> for i in range(group_size): ... b_fp8[i], b_scale[i] = per_block_cast_to_fp8(b_f32[i]) >>> >>> # Create group assignment >>> m_indices = torch.empty(m, device="cuda", dtype=torch.int32) >>> for i in range(group_size): ... row_slice = slice(i * m_per_group, (i + 1) * m_per_group) ... m_indices[row_slice] = i >>> >>> # Perform grouped GEMM >>> result = group_deepgemm_fp8_nt_groupwise( ... a_fp8, b_fp8, a_scale, b_scale, m_indices, out_dtype=torch.bfloat16 ... ) >>> print(result.shape) # torch.Size([256, 2048]) Notes ----- - This function requires NVIDIA Blackwell (SM100) architecture - The scaling factors should be generated using appropriate quantization functions like ``per_token_cast_to_fp8`` for `a` and ``per_block_cast_to_fp8`` for `b` - The function internally uses the DeepGEMM backend for optimized FP8 computation - All input tensors must be on the same CUDA device - The block size for scaling is determined by the ``scale_granularity_mnk`` parameter r) m_grouped_fp8_gemm_nt_contiguousNrr)r#r&rTrrrr;) rZr[r5r6r rr^rr&s r0group_deepgemm_fp8_nt_groupwiser'sjFEEEEE {/ k!'!*agaj !(SSS$$ G q'lC4I Jr2masked_m expected_mc ddlm} |M|p tj}tj|jd|jd|jd||j}| ||f||f||||S)Nr)0_check_m_grouped_fp8_gemm_nt_masked_problem_sizerr)r#r+rTrrrr;) rZr[r5r6r(r)rr^rr+s r0&_check_batch_deepgemm_fp8_nt_groupwiser,AsVUUUUU {/ k GAJ AGAJi    < ; G q'lC:?T  r2c ddlm} |M|p tj}tj|jd|jd|jd||j}| ||f||f|||||S)aPerform batch matrix multiplication with FP8 data types using DeepGEMM backend. This function performs a batch GEMM operation where each group in tensor `b` is multiplied with the corresponding group of rows in tensor `a`. The results of each group is masked by the `masked_m` tensor, which specifies which group each row belongs to. This is particularly useful for scenarios like mixture of experts (MoE) where different tokens are routed to different experts. The operation can be conceptualized as: >>> for i in range(num_groups): >>> output[i] = a[i][:masked_m[i]] @ b[i][:masked_m[i]].T Currently only supported on NVIDIA Blackwell (SM100) architecture. Parameters ---------- a : torch.Tensor Input tensor A of shape ``(batch_size, m, k)`` with FP8 data type (``torch.float8_e4m3fn``). Each slice ``a[i]`` represents a group of rows that will be multiplied with the corresponding group/expert in `b`. b : torch.Tensor Input tensor B of shape ``(batch_size, n, k)`` with FP8 data type (``torch.float8_e4m3fn``). Each slice ``b[i]`` represents a different group/expert that will be multiplied with the corresponding rows in `a`. a_scale : torch.Tensor Scaling factors for tensor `a` of shape ``(batch_size, m, k // block_size)`` with ``torch.float32`` dtype. These are typically generated from per-token quantization of the original float32 tensor. b_scale : torch.Tensor Scaling factors for tensor `b` of shape ``(batch_size, n // block_size, k // block_size)`` with ``torch.float32`` dtype. These are typically generated from per-block quantization of the original float32 tensor for each group. masked_m : torch.Tensor Masking tensor of shape ``(batch_size,)`` with ``torch.int32`` dtype. Each element specifies the effective rows to be multiplied in each group. For example, if ``masked_m[i] = j``, then first ``j`` rows in `a[i]` will be multiplied with group ``i`` in `b`. expected_m : int A value hint (which is a value on CPU) for the M expectation of each batch, correctly setting this value may lead to better performance. scale_granularity_mnk : Tuple[int, int, int], optional The granularity of the scaling factors as ``(m_granularity, n_granularity, k_granularity)``. Default is ``(1, 128, 128)`` which means per-token scaling for `a` and 128x128 block scaling for `b`. out : Optional[torch.Tensor], optional Pre-allocated output tensor of shape ``(batch_size, m, n)``. If not provided, a new tensor will be created. out_dtype : Optional[torch.dtype], optional Data type of the output tensor. If `out` is provided, this parameter is ignored. Default is ``torch.bfloat16``. Returns ------- torch.Tensor Output tensor of shape ``(batch_size, m, n)`` containing the results of the batch matrix multiplication. Examples -------- >>> import torch >>> from flashinfer.gemm import batch_deepgemm_fp8_nt_groupwise >>> from flashinfer.utils import per_token_cast_to_fp8, per_block_cast_to_fp8 >>> >>> # Setup: 2 groups, 128 tokens per group, 4096 hidden size, 2048 expert size >>> m, n, k = 128, 2048, 4096 >>> group_size = 2 >>> >>> # Create float32 inputs >>> a = torch.rand((group_size, m, k), device="cuda", dtype=torch.float32) >>> b = torch.rand((group_size, n, k), device="cuda", dtype=torch.float32) >>> masked_m = torch.randint(0, m, (group_size,), device="cuda", dtype=torch.int32) >>> a_fp8 = torch.empty_like(a, device="cuda", dtype=torch.float8_e4m3fn) >>> a_scale = torch.empty((group_size, m, k // 128), device="cuda", dtype=torch.float32) >>> b_fp8 = torch.empty_like(b, device="cuda", dtype=torch.float8_e4m3fn) >>> b_scale = torch.empty( ... (group_size, n // 128, k // 128), device="cuda", dtype=torch.float32 >>> ) >>> for i in range(group_size): ... a_fp8[i], a_scale[i] = per_token_cast_to_fp8(a[i]) ... b_fp8[i], b_scale[i] = per_block_cast_to_fp8(b[i]) >>> >>> expected_m = min(int(masked_m.float().mean()) + 1, m) >>> >>> # Perform batch GEMM >>> result = batch_deepgemm_fp8_nt_groupwise( ... a_fp8, b_fp8, a_scale, b_scale, masked_m, expected_m, out_dtype=torch.bfloat16 ... ) >>> print(result.shape) # torch.Size([2, 128, 2048]) Notes ----- - This function requires NVIDIA Blackwell (SM100) architecture - The scaling factors should be generated using appropriate quantization functions like ``per_token_cast_to_fp8`` for `a` and ``per_block_cast_to_fp8`` for `b` - The function internally uses the DeepGEMM backend for optimized FP8 computation - All input tensors must be on the same CUDA device - The block size for scaling is determined by the ``scale_granularity_mnk`` parameter r)m_grouped_fp8_gemm_nt_maskedNrr)r#r.rTrrrr;) rZr[r5r6r(r)rr^rr.s r0batch_deepgemm_fp8_nt_groupwiser/ZspBAAAAA {/ k GAJ AGAJi   ! G q'lC:?T Jr2ct}ddlm}t |jdz dz }||g|S)z4Get the FP8 block scale GEMM runner module for SM90.r )env nv_internal tensorrt_llm)r*r~jitr1rFLASHINFER_CSRC_DIRset_deepgemm_jit_include_dirsinit)r`jit_envdeepgemm_include_dirs r0#get_fp8_blockscale_gemm_runner_sm90r:su1 2 2 A A C CF$$$$$$#m3nD ((*>)?@@@ ;;==r2inputweight input_scale weight_scalecNt|jddgstd|jdkrtd|j|jdkrtd|j|j\}}|j\}} | |krtd|d| d } || zd krtd | d ||d zd krtd||jt jk} |jt jk} |jt jk} |jt jk}| r|rtd| rm|td|jt j krtd|j|j|jkrtd|jd|jn*| std|j|td| r|td||| zf}|| zdz | z|| zf}|j|k}|j|k}|s|std|d|d|j|jt j krtd|jn*|std|j|td||j||fkrtd|d |d!|j|j|jkrtd"|jd|j|jt jt j fvrtd#|j|%|j|krtd$|d|j|j}nW|p t j}|t jt j fvrtd#|t j ||||j%}t}| |||}d}|d kr;t j |t j|j%}|||||||||S)&u; Perform FP8 block-scaled GEMM with automatic swapAB optimization. This function automatically selects between normal and swapAB kernel based on the M dimension. For small M (< 32), it uses the swapAB kernel for better performance. Supported Dtype Combinations ----------------------------- - **BF16 + BF16 → BF16**: Both inputs BF16, internal quantization (no scales needed) - **BF16 + FP8 → BF16**: BF16 input, FP8 weight - **FP8 + FP8 → BF16** (W8A8): Both inputs FP8 with scales required Parameters ---------- input : torch.Tensor Input activation tensor of shape (M, K). - BF16 (torch.bfloat16) with internal quantization weight : torch.Tensor Weight tensor of shape (N, K). Can be: - FP8 (torch.float8_e4m3fn) with weight_scale required - BF16 (torch.bfloat16) for internal quantization input_scale : torch.Tensor, optional weight_scale : torch.Tensor, optional Scaling factors for weight. Required if weight is FP8. out : torch.Tensor, optional Output tensor of shape (M, N). If None, will be allocated. out_dtype : torch.dtype, optional Output data type. Default is torch.bfloat16. Returns ------- torch.Tensor Output tensor of shape (M, N) with dtype `out_dtype`. Examples -------- >>> import torch >>> from flashinfer.gemm import fp8_blockscale_gemm_sm90 >>> >>> M, N, K = 16, 4096, 4096 >>> device = "cuda" >>> >>> # BF16 inputs >>> input_bf16 = torch.randn(M, K, device=device, dtype=torch.bfloat16) >>> weight_bf16 = torch.randn(N, K, device=device, dtype=torch.bfloat16) >>> output = fp8_blockscale_gemm_sm90(input_bf16, weight_bf16) >>> print(output.shape) # torch.Size([16, 4096]) >>> >>> # Mixed: BF16 input + FP8 weight >>> from flashinfer.testing.utils import per_token_cast_to_fp8 >>> input_bf16 = torch.randn(M, K, device=device, dtype=torch.bfloat16) >>> weight_bf16 = torch.randn(N, K, device=device, dtype=torch.bfloat16) >>> weight_fp8, weight_scale = per_token_cast_to_fp8(weight_bf16) >>> output = fp8_blockscale_gemm_sm90(input_bf16, weight_fp8, None, weight_scale) >>> print(output.shape) # torch.Size([16, 4096]) >>> >>> # FP8 weight with 128x128 block scales >>> from flashinfer.testing.utils import per_block_cast_to_fp8 >>> weight_bf16 = torch.randn(N, K, device=device, dtype=torch.bfloat16) >>> weight_fp8, weight_scale = per_block_cast_to_fp8(weight_bf16) >>> # weight_scale has shape (N // 128, K // 128) >>> input_bf16 = torch.randn(M, K, device=device, dtype=torch.bfloat16) >>> output = fp8_blockscale_gemm_sm90(input_bf16, weight_fp8, None, weight_scale) >>> print(output.shape) # torch.Size([16, 4096]) Notes ----- - This function requires NVIDIA Hopper (SM90) architecture and CUDA 12.8+ - SwapAB kernel is automatically used when M < 32 (threshold) - For FP8 inputs, scaling factors must be provided - For BF16 inputs, quantization and scaling happen internally - Weight scales support two granularities: * Per-token (1x128 blocks): (N, K//128) * Per-block (128x128 blocks): (N//128, K//128) - Input scales only support per-token format: (M, K//128) - The function uses DeepGEMM backend with JIT compilation 9090azIfp8_blockscale_gemm_sm90 is only supported on SM90 (Hopper) architecture.r z#Input must be 2D (M, K), got shape z$Weight must be 2D (N, K), got shape z"K dimension mismatch: input has K=z, weight has K=rrz-K dimension must be divisible by block size (z ), got K=r z+N dimension must be divisible by 64, got N=zJFP8 input + BF16 weight is not supported (missing kernel implementation). Nz+input_scale is required when input is FP8. z!input_scale must be float32, got z&input_scale device mismatch. Expected rzMInput must be either FP8 (torch.float8_e4m3fn) or BF16 (torch.bfloat16), got zjinput_scale should not be provided for BF16 inputs. Use FP8 inputs if you want to provide external scales.z-weight_scale is required when weight is FP8. rz-weight_scale shape mismatch. Expected either z (per-token, 1x128 blocks) or z" (per-block, 128x128 blocks), got z"weight_scale must be float32, got zNWeight must be either FP8 (torch.float8_e4m3fn) or BF16 (torch.bfloat16), got zmweight_scale should not be provided for BF16 weights. Use FP8 weights if you want to provide external scales.z!Output shape mismatch. Expected (z, z), got rz:Output dtype must be torch.bfloat16 or torch.float16, got rr)rBr;rrrrrTrJrrrqrr:r,rIconfigure_workspacerun_gemm)r;r<r=r>r^rMrNK_weight BLOCK_SIZE input_is_fp8 weight_is_fp8 input_is_bf16weight_is_bf16expected_per_token_shapeexpected_per_block_shape is_per_token is_per_blockr2r2rDs r0fp8_blockscale_gemm_sm90rPsh U\D%= 9 9  W    zQLu{LLMMM {aN NNOOO ;DAq,KAx1}} M M M8 M M   J:~ TJ T TQR T T    2v{{JqJJKKK;%"55LLE$77MK5>1M\U^3N   X     JKK K   - -TARTTUU U   - -,,,"),,  .  %{%%   "I    LMM M$%qJ#7 %&^a%7J$FZ#X #)-EE #)-EE    I@XII0HII4@4FII    . .V,BTVVWW W / &|&&   #J    9A  NANNNN39NN  : % %TELTT TT  9U^U]; ; ;XSYXX   SY)%;%;O9OOCIOO I / U^U]; ; ;XYXX k!Qi EEE1 2 2F..q!Q77NIKek%,WWW ""9--- OOE63 \BBB Jr2rkrlrcd}d}d}||||z}||||z}||||||z} ||| fS)zDCalculate block scale dimensions using indestructible block formula.rrc||zdz |zSrrL)rZr[s r0div_upz+_calculate_block_scale_dims..div_upsA ar2rL) rkrlrrINDESTRUCTIBLE_128x4_BLOCK_M_NINDESTRUCTIBLE_128x4_BLOCK_KrSblock_scale_dim_mblock_scale_dim_nblock_scale_dim_ks r0_calculate_block_scale_dimsrYs&)"#$     q0114RR q0114RR vva$$&BCC & ' /1B BBr2c t|dkrtd|t|dkrtd||tjjtjjfvrtd||tjjtjjfvrtd||tjjtjjfvrtd||d} |d} |d } |d } t| | | |\} }}| | |f}| |z|df}| ||f}||zd|f}tjj }tj |}tj t|5\}}|d t!|t!|| }|d t!|t!|| }|d |||tjj}|d|||tjj}|||d|gd}|tjj||||dgd}|tjj|||tjjd}|tjj|d||}|t2jj|t2jj|t2jj|t2jj|t2jj| |!|"tj#j$tj#j%g|cdddS#1swxYwYdS)NrzA shape must be 3D, got zB shape must be 3D, got z)A type must be FP8_E4M3 or FP8_E5M2, got z)B type must be FP8_E4M3 or FP8_E5M2, got z&Output type must be BF16 or FP16, got rrr rZrr[rrrrrrrrT)&rXrr7rrrKrHrIrYrrTrUrrrrrZrrrrrrrrrrrrrrrrrrr rr) r r r;r rr<rrr;b_dimrkrrlrVrWrXrrrrrrrr!rrrrrrrrs r0'create_cudnn_execution_plans_mxfp8_gemmr\s 7||q=G==>>> 7||q=G==>>> eo.0HIIIMVMMNNN eo.0HIIIMVMMNNN eo.0DEEEJ&JJKKK AJE A A A $Aq!Z88<(*;/1BCO--  /1BCO-- )J Z & &v . .F &v.. / /C:E1g?? &   g?? &   (-||"# !3< (4( ( $(-||"# !3< (4( ( $!77  (: 8   &&u'<=== 77  ("A 8   &&u'<===<<  #o3     u4555 D!!//777'tz/000tz/000$,,T-E-KLLL$,,T-E-KLLL$$TZ%5666  ##%%% $$eo&79J%KLLLGCCCCCCCCCCCCCCCCCCsKQQ!$Q!rc t|j||j|t|jt|jt|||j }||dkr||n||S)N) r r r rr;r<rrr;rP) r\rrrLrr;r#r$r%)rZr[rr^rrQrs r0_get_cudnn_mxfp8_gemm_graphr^as 4217;;217;;29==x   E  || !!&))))  Lr2c bd}t||||||} t| |||||||dS)Nr)rZr[rr^rrQ)rrZr[rLrMrr_rQ)r^r3) rZr[rLrMrr^r_rQrrs r0_cudnn_gemm_mxfp8r`}slJ (      E# )      r2c<Gddt}|S)Nc eZdZdeejdedeefdZ d deejdede dejfd Z d S) 6_cudnn_gemm_mxfp8_runner..CudnnMxfp8GemmRunnerrGrHrIcdgSrKrLrMs r0rOzH_cudnn_gemm_mxfp8_runner..CudnnMxfp8GemmRunner.get_valid_tacticss 3Jr2rPFrQrRc P|\}}}}} } t||||| | j| || S)N)rZr[rLrMr^rr_rQ)r`rrUs r0raz>_cudnn_gemm_mxfp8_runner..CudnnMxfp8GemmRunner.forwardsP=C 9Aq'7C)9 !!)!1    Jr2NrbrcrLr2r0CudnnMxfp8GemmRunnerrcs & ) #Y     #(   &  !   \       r2rfrl)rfs r0_cudnn_gemm_mxfp8_runnerrgs=}<  ! !!r2cg}d|vr!|t|s Jdtj}||||||g} |d|t | \} } | | | dS)Nr7r( mxfp8_gemmr))rrgr r,r-r9r:s r0mxfp8_gemm_sm100rjsG,/11222 /////7 MOOEGWc+; rs- !!!!!!111111111111FFFFFF ((((((                    ,+++++&&&&&&,,,,,,,,,,,,888888888888888888999999111111000000000000******:::::: LLLOO   DA ISEE_EEEEEN  211111*-T)%el%S %%%%  QQQhsCj))#'"^#'16 | | %, {  5<   - .*).sCj))#'"^#'16   |  |  %,  {  5<   - .   *) &$("&"^16 | | 5<    %,   { - .<$("&"^16Cy | | 5<    %,  {- .,/$-* $("&"^16cc |c |c 5< c  c %,  c { c- .c \ccccLsCj))#'"^"+   |  |  %,  {  Y    *) #'"^"+  | | %, {  Y  .#'"^"+  Cy  |  |  %,   {  Y     0.+ #'"^"+ EE |E |E %, E{ E Y  E  \ EEEEP     fffR  $$$N!-    2 $      ( (  !!!& !!!H".   2 $      ( (  """&) |) |) ,)  )  ) l )s)) ))))>) |) |)\)\ )  ) l )s)) ))))@/S/c////d  FFFFF U11 ;16:1111h "& c0c0 |c0 |c0 ,c0  c0 %,  c0  \ c0c0c0c0L <<<L.2;; |; \; |; ;  ;  ;U\*;;;;J.2;; |; \; |; ;  ;  ;U\*;;;;|~~~~~~~~B     4      B    ej/ ek     u{     aaaH  (( ((((h" "  " " " " Z         F MMM`@@@, 8u{ 8 8 8 8| | |\  \  %,  [8   03330AAA:%)"^"&-A kk |k |k EL !k{ k %,  k ) * kkkkf%)"^"&,, |, |,|,| , EL ! , { , %, ,,, ,,,,h%)"^"&%) | |||  EL !  {  %, l @G G G ^%)"^"&#=C55 |5 |5|5| 5 EL ! 5 { 5 %, 5559 :55555p77788 %)"^"&#=C99 |9 |9|9| 9 EL ! 9 { 9 %, 9999 :9999989xsCj)) %)"^"&#=C | |||  EL !  {  %, 9 :*).77788 %)"^"&#=C | |||  EL !  {  %, 9 :984%)"^"&#=D$E$ECy$E |$E |$E| $E | $E EL ! $E{$E %, $E$E$E9 :$E$E$E$E$EN""")L   2 $     3 3   ' '  0+l   2 $     5 5   ' '  0,.0 ,) %)"^"&#=CII |I |I|I| I EL ! I { I %, III9 :II \IIIIX:::;;#'=E   |  | \ \  ;  %,  9 :   <; :::;;#'=E   |  | \ \  ;  %,  9 :   <; 77788#'=E   |  | \ \  ;  %,  9 :   98 (#'=E   |  | \ \  ;  %,  9 :    (#'=ECy | |\  \  ;  %, 9 :<+-/ -*#'=EXX |X |X\X\ X ; X %,  X9 :X \XXXXv22233 6:2?"&'+,5 | |\\  wy12   !c3/ %,  $( )43$sCj)) 6:2?"&'+,5 | |\\  wy12   !c3/ %,  $( )*)26:2?"&'+,5 | |\\  wy12   !c3/ %,  $( )8=;;  6:2?"&'+,5JJ |J |J\J\ J wy12 J  J!c3/J %, J $J( )J \JJJJZ QQQh22233 6:"&'+$$ |$ |$\$\ $ wy12 $  $ %, $ $$$$43$N< 6:"&'+ | |\\  wy12    %,  $ \  8222333@+/"&'+AA |A |A\A\ A l A !c3/ Ai(A A %, A $AAA43AH@3@+/"&'+uu |u |u\u\ u l u !c3/ ui(u u %, u $u \uuu  up77788"&'+OO |O |O\O\ O l O  O O O OO %, O $OOO98OdH"&'+ll |l |l\l\ l l l  l l l ll %, l $l \lll  l`!D*l****.  sCj))3@"&'+ | |\\  |  !c3/  %,  $ *).D3@"&'+zz |z |z\z\ z | z !c3/ z %, z $zzz  zzsCj))3@"&'+ | |\\  l   !c3/ %,  $ *)073@"&'+ | |\\  l   !c3/ %,  $  D    +/+/"&'+ VV <V LV%,'V5<( V %,  V  $ V \VVVVrC CCC(+C 3S=CCCC0 }}}F#^"&  | |{ %,      B#^"&%)   |  | | |  {  %,  l      F"""D) |) |)\)\ )  ) l )s)) ))))4sCj))#' '   |  | \ \  ;  %,  W    *)       #' ' | |\\  ;  %,  W :#' '  Cy  |  | \  \  ;  %,  W      -/, #' ':: |: |:\:\ : ; : %,  :W : \::::::s,B33C=:C=>5C88C=