A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 21:00 UTC
[Hardware][Intel-Gaudi] Regional compilation support (#13213)
Kacper Pietkun committed
Feb 28, 2025 at 09:51 UTC
b91660ddb8fe5a89bb14f8aa9943caf9306e7288
1 file changed
+37
-6
vllm/worker/hpu_model_runner.py
+37
-6
@@ -39,7 +39,10 @@ from vllm.lora.layers import LoRAMapping
39
from vllm.lora.request import LoRARequest
40
from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
41
from vllm.model_executor import SamplingMetadata
42
+from vllm.model_executor.layers.layernorm import RMSNorm
43
from vllm.model_executor.layers.sampler import SamplerOutput
44
+from vllm.model_executor.layers.vocab_parallel_embedding import (
45
+ VocabParallelEmbedding)
46
from vllm.model_executor.model_loader import get_model
47
from vllm.multimodal import (MULTIMODAL_REGISTRY, BatchedTensorInputs,
48
MultiModalKwargs)
@@ -311,10 +314,38 @@ class HpuModelAdapter:
314
self.block_size = vllm_config.cache_config.block_size
315
self.dtype = vllm_config.model_config.dtype
316
enforce_eager = vllm_config.model_config.enforce_eager
317
+
318
if not htorch.utils.internal.is_lazy() and not enforce_eager:
315
- self.model = torch.compile(self.model,
316
- backend='hpu_backend',
317
- dynamic=False)
319
+ if os.getenv('VLLM_REGIONAL_COMPILATION',
320
+ 'true').lower() == 'true':
321
+ self.regional_compilation_layers_list = [
322
+ RMSNorm, VocabParallelEmbedding
323
+ ]
324
+ self._regional_compilation(self.model)
325
+ else:
326
+ self.model = torch.compile(self.model,
327
+ backend='hpu_backend',
328
+ dynamic=False)
329
+
330
+ def _regional_compilation(self,
331
+ module,
332
+ parent_module=None,
333
+ module_name=None):
334
+ if isinstance(module, torch.nn.ModuleList):
335
+ for children_name, children_module in module.named_children():
336
+ self._compile_region(module, children_name, children_module)
337
+ elif any(
338
+ isinstance(module, layer)
339
+ for layer in self.regional_compilation_layers_list):
340
+ self._compile_region(parent_module, module_name, module)
341
+ else:
342
+ for children_name, children_module in module.named_children():
343
+ self._regional_compilation(children_module, module,
344
+ children_name)
345
+
346
+ def _compile_region(self, model, name, module):
347
+ module = torch.compile(module, backend='hpu_backend', dynamic=False)
348
+ setattr(model, name, module)
349
350
def _set_attn_bias(self, attn_metadata, batch_size, seq_len, device,
351
dtype):
@@ -1575,9 +1606,9 @@ class HPUModelRunnerBase(ModelRunnerBase[TModelInputForHPU]):
1606
list(sorted(self.bucketing_global_state.decode_buckets)))
1607
1608
if not htorch.utils.internal.is_lazy() and not self.enforce_eager:
1578
- cache_size_limit = len(
1579
- self.bucketing_global_state.prompt_buckets) + len(
1580
- self.bucketing_global_state.decode_buckets) + 1
1609
+ cache_size_limit = 1 + 3 * (
1610
+ len(self.bucketing_global_state.prompt_buckets) +
1611
+ len(self.bucketing_global_state.decode_buckets))
1612
torch._dynamo.config.cache_size_limit = max(
1613
cache_size_limit, torch._dynamo.config.cache_size_limit)
1614
# Multiply by 8 to follow the original default ratio between