
    ;il                     l	   S SK Jr  S SKrS SKrS SKJr  S SKrS SKJ	r	J
r
  SSKJr  SSKJrJr  SS	KJr  \R$                  " \5      r\R*                  " \R,                  5      R.                  \R*                  " S
5      R.                  :  a)  \" SS5      S\R0                  S\R0                  4S j5       rS\R0                  S\R0                  S\S\S\R8                  S\R0                  SS4S jrS\R0                  S\R0                  S\R0                  S\S\R8                  S\R0                  SS4S jrS\R0                  S\R0                  S\\   S\R0                  S\R0                  S\S\R0                  SS4S jr\ " \\5      (       Gd  \RC                  S5        \(       aj  SSK"J#r$  \" SS5      " \$RJ                  5        \" SS5      " \$RL                  5        \" SS5      " \$RN                  5        \" S S5      " \$RP                  5        \" S!S5      S\R0                  S\R0                  S\S\S"\\   S\R8                  S\R0                  4S# j5       r\" S$S5      S\R0                  S\R0                  S\R0                  S\S\R8                  S\R0                  4S% j5       r\" S&S5      S\R0                  S\R0                  S\R0                  S\S\R8                  S\R0                  SS4S' j5       r\" S(S5      S\R0                  S\R0                  S\\   S\R0                  S\R0                  S\S\R0                  4S) j5       r\" S*S5      S\R0                  S\R0                  S\\   S\R0                  S\R0                  S\S\R0                  SS4S+ j5       rg\(       a  \RC                  S,5        SSK"J#r$  \" SS5      " \$RJ                  5        \" S&S5      " \$RR                  5        \" S$S5      " \$RT                  5        \" SS5      " \$RL                  5        \" S-S5      " \$RV                  5        \" S!S5      " \$RX                  5        \" S(S5      " \$RZ                  5        \" SS5      " \$RN                  5        \" S S5      " \$RP                  5        g\R]                  S.5        g)/    )SequenceN)version)_get_tensor_streamget_ptr   )register_kernel) ErrorHandlerMockBNBNativeLibrarylib   )triton_availablez2.9z bitsandbytes::int8_linear_matmulxpuABc                     [         R                  " U R                  SU R                  S   5      UR	                  5       5      R                  " / U R                  S S QUR                  S   P76 $ )Nr   )torch_int_mmreshapeshapet)r   r   s     m/home/dmtnaga/Documents/work/airagagent/rag_env/lib/python3.13/site-packages/bitsandbytes/backends/xpu/ops.py_r      si    }}IIb!''"+&CCE
 ' - 773B<- "#- 	-    absmax	blocksize
quant_typedtypeoutreturnc           	      V   S [        U 5      [        U5      [        U5      [        R                  " U5      [        R                  " UR                  5       5      [	        U 5      4nU[
        R                  :X  a.  US:X  a  [        R                  " U6   g [        R                  " U6   g U[
        R                  :X  a.  US:X  a  [        R                  " U6   g [        R                  " U6   g U[
        R                  :X  a.  US:X  a  [        R                  " U6   g [        R                  " U6   g g )Nfp4)r   ctc_intnumelr   r   bfloat16r
   cdequantize_blockwise_bf16_fp4cdequantize_blockwise_bf16_nf4float16cdequantize_blockwise_fp16_fp4cdequantize_blockwise_fp16_nf4float32cdequantize_blockwise_fp32_fp4cdequantize_blockwise_fp32_nf4)r   r   r   r   r   r   argss          r   _dequantize_4bit_implr/      s     	


1D ..5..5	%--	..5..5	%--	..5..5	 
 r   codec           	         [        U5      [        U 5      [        U5      [        U5      [        R                  " U5      [        R                  " U R                  5       5      [	        U 5      4nU[
        R                  :X  a  [        R                  " U6   g U[
        R                  :X  a  [        R                  " U6   g U[
        R                  :X  a  [        R                  " U6   g g )N)r   r"   r#   r$   r   r   r(   r
   cdequantize_blockwise_fp16r%   cdequantize_blockwise_bf16r+   cdequantize_blockwise_fp32)r   r   r0   r   r   r   r.   s          r   _dequantize_blockwise_implr5   =   s     	


1D &&-	%..	 &&-	%--	&&- 
 r   shapeBc                    [         R                  " S5      n[         R                  " US   5      n[         R                  " US   5      n	Un
[         R                  " U R                  S   S-   S-  5      nUn[        U 5      nU R                  [
        R                  :X  ad  [        R                  " UUU	[        U 5      [        U5      [        U5      [        U5      [        U5      U
UU[         R                  " U5      U5        g U R                  [
        R                  :X  ad  [        R                  " UUU	[        U 5      [        U5      [        U5      [        U5      [        U5      U
UU[         R                  " U5      U5        g U R                  [
        R                  :X  ad  [        R                  " UUU	[        U 5      [        U5      [        U5      [        U5      [        U5      U
UU[         R                  " U5      U5        g g )N   r   r   r   )r"   c_int32r   r   r   r   r(   r
   cgemv_4bit_inference_fp16r   r%   cgemv_4bit_inference_bf16r+   cgemv_4bit_inference_fp32)r   r   r6   r   r0   r   r   mnkldaldbldcstreams                 r   _gemv_4bit_implrD   Q   s    	

1A


6!9A


6!9A
C
**aggbkAo!+
,C
C"Fww%--%%AJAJFODMCLJJy!	
 
ENN	"%%AJAJFODMCLJJy!	
 
EMM	!%%AJAJFODMCLJJy!	
 
"r   z*Register sycl bitsandbytes kernels for XPU)opsz bitsandbytes::quantize_blockwisezbitsandbytes::quantize_4bitz-bitsandbytes::optimizer_update_8bit_blockwisez$bitsandbytes::optimizer_update_32bitzbitsandbytes::dequantize_4bitr   c           	      Z    [         R                  " XEU R                  S9n[        XX#XVS9  U$ )N)r   devicer   )r   emptyrG   r/   )r   r   r   r   r   r   r   s          r   r   r      s)     kk%QXX>aO
r   z"bitsandbytes::dequantize_blockwisec           	      D    [         R                  " XS9n[        XX#XES9  U$ )Nr   rH   )r   
empty_liker5   r   r   r0   r   r   r   s         r   r   r      s%     q."1duN
r   z&bitsandbytes::dequantize_blockwise.outc           	         ^ ^^ [         R                  " TR                  T:H  UU4S j5        [         R                  " TR                  T R                  :H  U U4S j5        [	        T XUTTS9  g )Nc                  (   > ST  STR                    3$ NzExpected out.dtype == , got rK   )r   r   s   r   <lambda>_.<locals>.<lambda>   s    3I%PVWZW`W`Va1br   c                  <   > ST R                    STR                    3$ )NExpected out.shape == rQ   r   r   r   s   r   rR   rS           5KAGG9TZ[^[d[dZe3fr   rH   )r   _checkr   r   r5   rM   s   `   ``r   r   r      sI     	SYY%')bcSYY!'')+fg"1fIu#Nr   zbitsandbytes::gemv_4bitc           
          / U R                   S S QUS   P7n[        R                  " X`R                  U R                  S9n[        XX#XEUS9  U$ )Nr   r   )rG   r   rH   )r   r   rI   rG   r   rD   )r   r   r6   r   r0   r   r   r   s           r   r   r      sL     +!''#2,*q	*kk%@fd3G
r   zbitsandbytes::gemv_4bit.outc           
      
  ^ ^^ [         R                  " TR                  / T R                  S S QTS   P7:H  U UU4S j5        [         R                  " TR                  T R                  :H  U U4S j5        [	        T UTX4UTS9  g )Nr   r   c                  R   > S/ T R                   S S QTS   P7 STR                    3$ )NrU   r   r   rQ   rV   )r   r   r6   s   r   rR   rS      s1    ,-Gqwws|-GVAY-G,HsyykZr   c                  <   > ST R                    STR                    3$ rP   rK   rW   s   r   rR   rS      rX   r   rH   )r   rY   r   r   rD   )r   r   r6   r   r0   r   r   s   ` `   `r   r   r      si     	II31773B<333Z	
 	SYY!'')+fg1ffI3Gr   z,Register triton bitsandbytes kernels for XPUz!bitsandbytes::dequantize_4bit.outzaRegister pytorch bitsandbytes kernels for XPU because no native library or triton packages found.)/collections.abcr   ctypesr"   logging	packagingr   r   bitsandbytes.functionalr   r   _opsr   
cextensionr	   r
   utilsr   	getLogger__name__loggerparse__version__releaseTensorr   intstrr   r/   r5   rD   
isinstanceinfotritonrE   
triton_opsquantize_blockwisequantize_4bitoptimizer_update_8bit_blockwiseoptimizer_update_32bitdequantize_blockwise_inplacedequantize_blockwisedequantize_4bit_inplacedequantize_4bit	gemv_4bitwarning r   r   <module>r~      sh   $     ? # ? $			8	$ ==""#++w}}U/C/K/KK7?-U\\ -ell - @-6||6LL6 6 	6
 ;;6 
6 
6D.||."\\.16.JM.V[VaVa.hmhtht.	.(A
||A
||A
 SMA
 LL	A

 ,,A
 A
 
A
 
A
J #788
KK<= .:EB:C`C`a5u=j>V>VWGO66	
 	>FzGhGhi4e<
<<

 
 	

 }
 {{
 

 =
 95A<<!&5:\\NQZ_ZeZe	 B =uE
O<<
O
O ll
O 	
O
 {{
O \\
O 

O F
O .6<<<<  	
 ll  
 7 2E:H<<H<<H H 	H
 llH H \\H 
H ;H 
KK>?*6>z?\?\]<eDZElElm8%@A`A`a159*:R:RS7?
@b@bc3U;J<V<VW-u5j6J6JKCUKJLvLvw:EB:CdCde
NNvwr   