-`iz4ddlZddlZddlmZddlmZmZddlmZm Z ddl m Z ddl m Z mZmZddlZddlZddlmZddlmZmZmZddlmZdd lmZdd lmZeeZed Z ed Z!d ej"j#j$j%de de ddfdZ&d ej"j#j$j%de de ddfdZ'ededfdZ(GddZ)dS)N)abstractmethod)Callable Generator)contextmanager nullcontext)CodeType)Any ParamSpecTypeVar)CompilationMode CUDAGraphModeget_current_vllm_config)DynamicShapesType) init_logger)layerwise_nvtx_marker_contextRPselfargskwargsreturncdS)z-No-op to skip the GLOBAL_STATE guard entirelyNrrrs l/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/vllm/compilation/wrapper.py_noop_add_global_state_guardr  DcdS)z:No-op to skip the TORCH_FUNCTION_MODE_STACK guard entirelyNrrs r)_noop_add_torch_function_mode_stack_guardr "rrNNNc#Ktjjjjj}tjjjjj}tjjj}tjjj } dtjj_dtjj_ ttjjjj_ttjjjj_dV|tjjjj_|tjjjj_|tjj_|tjj_ dS#|tjjjj_|tjjjj_|tjj_|tjj_ wxYw)aContext manager for compilation settings and patches. This manager: 1. Sets higher dynamo cache limits for compilation. (Needed for qwen2_5_vl see test_qwen2_5_vl_evs_functionality). Generally a recompilation can happen whenever we use a new backend instance in torch.compile. 2. Patches out add_global_state_guard to skip GLOBAL_STATE guards 3. Patches out add_torch_function_mode_stack_guard to skip TORCH_FUNCTION_MODE_STACK guards. 4. Restores everything when compilation completes ii N) torch_C_dynamoguards GuardManageradd_global_state_guard#add_torch_function_mode_stack_guardconfigcache_size_limitaccumulated_cache_size_limitrr )original_global_state_guard(original_torch_function_mode_stack_guardoriginal_cache_sizeoriginal_accumulated_caches r_compilation_contextr1)s\ ,C  ,P- -.?!&!5!RW04 -<@ 9 ) ,C 6 ,P  ( ,C 5 ,P1D -TorchCompileWithNoGuardsWrapper._call_with_optional_nvtx_rangeis  . T Iv,,K.299$.:QRR#"  : ([$9&99  : : : : : : : : : : : : : : :: {D+F+++sA00A47A4Ncd|_t}||_|jj}|jj|_|td|j |}i}t|tr|dkr |jj }d|_ |jjj|_|jjj}|t"jkrP|jrCt&jr Jdt&jr|t,jks Jdd|d<nd |d<|j}|t,jkr.t&jr Jd |jr Jd |j}t7}t&jrkt9t:jjd r&t:jj d }n&d}|dz }|dz }tB"||5t;j#|dd|||_$dddn #1swxYwYt&jrB|t"jkr4t:jj%&|j'd|_(dSdSdS)NFz)Compilation mode cannot be NO_COMPILATIONinductorTz\compilation_config.dynamic_shapes_config.evaluate_guards requires VLLM_USE_BYTECODE_HOOK=0. zGevaluate_guards for backed shapes requires VLLM_USE_AOT_COMPILE=False. cd|DS)Nc$g|] }|jdkS) SHAPE_ENV) guard_type).0entrys r zNTorchCompileWithNoGuardsWrapper.__init__....s-8888=E$ 3888rrxs rz:TorchCompileWithNoGuardsWrapper.__init__..s 88AB888rguard_filter_fncd|DS)Ncg|]}dS)Fr)rM_s rrOzNTorchCompileWithNoGuardsWrapper.__init__....s7I7I7I!7I7I7IrrrPs rrRz:TorchCompileWithNoGuardsWrapper.__init__..s7I7Iq7I7I7Irz;UNBACKED dynamic shapes requires VLLM_USE_BYTECODE_HOOK=0. z)UNBACKED dynamic shapes do not add guardsenable_aot_compile)rWz/torch._dynamo.config.enable_aot_compile is not z.available. AOT compile is disabled and please z+upgrade PyTorch version to use AOT compile.) fullgraphdynamicbackendoptions))compiledr vllm_configcompilation_configmodeobservability_configenable_layerwise_nvtx_tracingr< RuntimeError init_backend isinstancestrinductor_compile_config first_compiledynamic_shapes_configevaluate_guardstyper STOCK_TORCH_COMPILEenvsVLLM_USE_BYTECODE_HOOKVLLM_USE_AOT_COMPILErBACKEDr7UNBACKEDr8rr6r#r%r*patchloggerwarningcompile_compiled_callable convert_frameregister_bytecode_hook bytecode_hook_compiled_bytecode) rr]r_rZr[ds_type compiled_ptr aot_contextmsgs r__init__z(TorchCompileWithNoGuardsWrapper.__init__ys -// &-2  , J + <JKK K0==kJJ gs # # M:(=(=!4LG!  * @ P 0FK ?6 6 6# J6:6 ,#&7&>>>>7?>> ..)**.J-I)* L  '0 0 0 2  M  2+ X X-X X X+ >t?Q R R R7;D # # # < <+V+Vs(HHHct|jds tddtjdz|j||fS)N aot_compilezuaot_compile is not supported by the current configuration. Please make sure torch.compile is enabled with the latest z)version of PyTorch (current using torch: ))r6rurbr# __version__rrs rrz+TorchCompileWithNoGuardsWrapper.aot_compilesct. >> ORe>ORRRS  &22D&>BBBrctjr|jjjt jkr |j|i|S|jsLtj j | |j|jg|Ri|S|5|j|jg|Ri|cdddS#1swxYwYdS|js|jst'ntjd}d|_t-5|5|j|jg|Ri|cdddcdddS#1swxYwYddddS#1swxYwYdS)Nfail_on_recompileF)rlrmr]r^r_r rkruryr#r% eval_frameremove_from_cacheoriginal_code_objectrF_dispatch_to_compiled_coder7rgrircompiler set_stancer1)rrrrEs r__call__z(TorchCompileWithNoGuardsWrapper.__call__s  &  38"677/t.????*  (::4;T;T;V;VWWW:t:+.26<4466>4> '+/5%D-1-AD ^../BCC  "'D %''    :t:+.26<                                         sH CC C E*#E9 E*E E*E E**E.1E.cdS)Nrrs rr7z'TorchCompileWithNoGuardsWrapper.forwards9<rc$|jjjS)z6Return the original code object of the forward method.)r@r7__code__)rs rrz4TorchCompileWithNoGuardsWrapper.original_code_objects~%..rold_codenew_codec||urdStj}|rd|jr]|j}|jj}|jjtj j d}|dkr|dkrn |r|j]|j d}|j|ksJ|j d|urdS||_ |j }|r|dz }|s ddl}||} t%|d 5} | | dddn #1swxYwYt(d |n#t,$rYnwxYw|j jjt2jkr6d |jvr/ddl}||} d | } t9| dSdS) z8Hook to save the compiled bytecode for direct execution.N_compilezconvert_frame.pyframerztransformed_code.pyrwz#Dynamo transformed code saved to %supdateaBAssigning / modifying buffers of nn.Module during forward pass is not allowed when using cudagraph inside the compiler because it will cause silent errors. Please use eager mode or fix the code. The following code contains clues about which buffer is being modified (please search for the usage of the function `update`): )rsys _getframef_backf_codeco_name co_filenamesplitospathsepf_localsryr]compile_debug_dump_pathexistsdepyf decompileopenwriterrdebug Exceptionr^cudagraph_moder NONEco_namesrb) rrrr code_name file_namerdecompiled_filersrcfr}s rrxz-TorchCompileWithNoGuardsWrapper.bytecode_hooks] 44466 6 6 F    LE ,I 066rw{CCBGIJ&&98J+J+J     w'|x'''' >& ! - - F"*7799  "%::O"))++  !LLL//(33Cos33%q %%%%%%%%%%%%%%%LL!FXXXX D   / >-BT T TH--- LLL//(++CRMP RR s## # U T--s69)E+"E8 E+EE+ E E++ E87E8r!c#K|}|jJ|j|jj_ dV||jj_dS#||jj_wxYw)a+ Context manager to dispatch to internally compiled code for torch<2.8. Why does this work? Because Dynamo guarantees that the compiled bytecode has exactly the same arguments, cell variables, and free variables as the original code. Therefore we can directly switch the code object in the function and call it. See https://dev-discuss.pytorch.org/t/what-is-the-relationship-requirement-among-original-bytecode-transformed-bytecode-and-bytecode-returned-by-hooks-in-dynamo/1693/7 for more details. N)rryr@r7r)roriginals rrz:TorchCompileWithNoGuardsWrapper._dispatch_to_compiled_code7sr,,..&222*.*A' 7 EEE.6DN " + + +hDN " + 6 6 6 6s AA!)rN)rA __module__ __qualname____doc__r r8rrrrrrFr~rrrr7rrrxrrrrrrr3r3Zs-#------ ,#AqD>,23&,DEH, ,,,, T<T<T<Trs ////////22222222********** OOOOOOOOOO555555######GGGGGG X   GCLL IcNN (  ! . 7: FI       (  ! . 7: FI      -Wi(89-W-W-W-W`o7o7o7o7o7o7o7o7o7o7r