Reduce number of compilations when dynamic shapes is used (#90)

It fixes the issue with to many compilation for Pytorch dynamic shapes (
when VLLM_T_COMPILE_DYNAMIC_SHAPES=1)
It allows making dynamic shapes for registered buffers (see
UnspecializedParamBufferSource in PyTorch) by setting dynamo config.
It also enables dynamic_shapes_compilation by default.

---------

Signed-off-by: Andrzej Kotłowski <[email protected]>

Author: anko-intel

vllm_gaudi/extension/features.py

@@ -78,7 +78,7 @@ def get_features():
         ValueFromList('bucketing_strategy', bucketing_strategies),
         Value('defrag', False),
         Value('regional_compilation', True, env_var='VLLM_T_COMPILE_REGIONAL_COMPILATION', env_var_type=boolean),
-        Value('dynamic_shapes_compilation', False, env_var='VLLM_T_COMPILE_DYNAMIC_SHAPES', env_var_type=boolean),
+        Value('dynamic_shapes_compilation', True, env_var='VLLM_T_COMPILE_DYNAMIC_SHAPES', env_var_type=boolean),
         Value('fullgraph_compilation', False, env_var='VLLM_T_COMPILE_FULLGRAPH', env_var_type=boolean),
     ]
     return split_values_and_flags(features)

vllm_gaudi/v1/worker/hpu_model_runner.py

@@ -2337,6 +2337,9 @@ def _maybe_compile(self, *args, **kwargs):
         """Entrypoint for a torch.compilation of the model"""
         if (not is_fake_hpu() and not htorch.utils.internal.is_lazy()
                 and not self.vllm_config.model_config.enforce_eager):
+            # force_parameter_static_shapes = False  alows to use dynamic
+            # shapes on tensors added to module via register_buffer()
+            torch._dynamo.config.force_parameter_static_shapes = False
             self.compile_config = HPUCompileConfig()
             if self.compile_config.regional_compilation:
                 self._compile_methods()