Merge branch 'habana_main' into adobrzyniewicz/tflops_main

.github/workflows/clang-format.yml

@@ -2,13 +2,13 @@ name: clang-format
 
 on:
  # Trigger the workflow on push or pull request,
- # but only for the main branch
+ # but only for the habana_main branch
  push:
  branches:
- - main
+ - habana_main
  pull_request:
  branches:
- - main
+ - habana_main
 
 jobs:
  clang-format:

.github/workflows/mypy.yaml

@@ -2,13 +2,13 @@ name: mypy
 
 on:
  # Trigger the workflow on push or pull request,
- # but only for the main branch
+ # but only for the habana_main branch
  push:
  branches:
- - main
+ - habana_main
  pull_request:
  branches:
- - main
+ - habana_main
 
 jobs:
  ruff:
@@ -50,4 +50,6 @@ jobs:
  mypy vllm/transformers_utils --config-file pyproject.toml
  mypy vllm/usage --config-file pyproject.toml
  mypy vllm/worker --config-file pyproject.toml
+ mypy vllm/hpu --config-file pyproject.toml
+
 

.github/workflows/ruff.yml

@@ -2,13 +2,13 @@ name: ruff
 
 on:
  # Trigger the workflow on push or pull request,
- # but only for the main branch
+ # but only for the habana_main branch
  push:
  branches:
- - main
+ - habana_main
  pull_request:
  branches:
- - main
+ - habana_main
 
 jobs:
  ruff:

.github/workflows/yapf.yml

@@ -2,13 +2,13 @@ name: yapf
 
 on:
  # Trigger the workflow on push or pull request,
- # but only for the main branch
+ # but only for the habana_main branch
  push:
  branches:
- - main
+ - habana_main
  pull_request:
  branches:
- - main
+ - habana_main
 jobs:
  yapf:
  runs-on: ubuntu-latest

format.sh

@@ -113,6 +113,7 @@ mypy vllm/spec_decode --config-file pyproject.toml
 mypy vllm/transformers_utils --config-file pyproject.toml
 mypy vllm/usage --config-file pyproject.toml
 mypy vllm/worker --config-file pyproject.toml
+mypy vllm/hpu --config-file pyproject.toml
 
 
 # If git diff returns a file that is in the skip list, the file may be checked anyway:

vllm/hpu/ops.py

@@ -15,15 +15,18 @@
 
 import vllm.hpu.utils as hpu_utils
 from vllm.worker.profiler import Profiler
+from vllm.logger import init_logger
 
-PA_SPLIT_VALUE = (os.environ.get('PA_SPLIT_VALUE', '1') == '1')
-
+logger = init_logger()
+HPUFusedRMSNorm = None
+try:
+ from habana_frameworks.torch.hpex.normalization import FusedRMSNorm
+ HPUFusedRMSNorm = FusedRMSNorm
+except ImportError:
+ logger.warning("Could not import HPU FusedRMSNorm kernel. "
+ "vLLM will use forward_native implementation of RMSNorm.")
 
-def silu_and_mul(output, input):
- d = input.shape[-1] // 2
- silu = torch.nn.SiLU().to(input.device)
- x, y = torch.split(input, d, dim=-1)
- output.copy_(silu(x) * y)
+PA_SPLIT_VALUE = (os.environ.get('PA_SPLIT_VALUE', '1') == '1')
 
 
 def fetch_from_cache(cache, blocks, permutations):
@@ -66,8 +69,7 @@ def paged_attention_v1(query,
  keys = [k.unflatten(1, (kv_heads, 1)) for k in keys]
  mask = mask.unsqueeze(2)
 
- attn_weights = [torch.matmul(query, k) for k in keys]
- attn_weights = torch.cat(attn_weights, dim=-1)
+ attn_weights = torch.cat([torch.matmul(query, k) for k in keys], dim=-1)
  if alibi_slopes is not None:
  attn_weights.add_(alibi_slopes[:, :, -attn_weights.size(2):,
  -attn_weights.size(3):])
@@ -103,12 +105,9 @@ def paged_attention_v1(query,
  return attn_weights.squeeze(-2)
 
 
-def silu_and_mul_wrapper(x: torch.Tensor) -> torch.Tensor:
+def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
  d = x.shape[-1] // 2
- output_shape = (x.shape[:-1] + (d, ))
- out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
- silu_and_mul(out, x)
- return out
+ return F.silu(x[..., :d]) * x[..., d:]
 
 
 def static_fused_moe(hidden_states, w1, w2, score, topk):
@@ -133,13 +132,10 @@ def static_fused_moe(hidden_states, w1, w2, score, topk):
  htorch.core.mark_step()
 
  for expert_idx in range(num_experts):
- padded_weight = padded_weights[expert_idx]
- current_state_static = hidden_states.reshape(-1, D)
- w_output = silu_and_mul_wrapper(
- torch.matmul(current_state_static, w1[expert_idx].transpose(0, 1)))
+ w_output = torch.matmul(hidden_states, w1[expert_idx].transpose(0, 1))
+ w_output = silu_and_mul(w_output)
  w_output = torch.matmul(w_output, w2[expert_idx].transpose(0, 1))
- current_hidden_states_static = w_output * padded_weight
- final_hidden_states += current_hidden_states_static
+ final_hidden_states += w_output * padded_weights[expert_idx]
  htorch.core.mark_step()
 
  return final_hidden_states.view(-1, D)
@@ -166,7 +162,8 @@ def prompt_attention(
  query = query.unflatten(1, (kv_heads, -1))
  key = key.unflatten(1, (kv_heads, 1))
  value = value.unflatten(1, (kv_heads, 1))
- attn_bias = attn_bias.unsqueeze(2)
+ if attn_bias is not None:
+ attn_bias = attn_bias.unsqueeze(2)
  attn_weights = torch.matmul(query * scale, key.transpose(-1, -2))
  if attn_bias is not None:
  attn_weights.add_(attn_bias)

vllm/hpu/rotary_embed.py

@@ -20,7 +20,7 @@
  except ImportError:
  logger.warning("Could not import HPU FusedRoPE kernel. "
  "vLLM will use forward_native implementation of RoPE.")
- FusedRoPE = None
+  FusedRoPE = None
 else:
  FusedRoPE = None
 

vllm/model_executor/layers/layernorm.py

@@ -6,19 +6,8 @@
 
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
-from vllm.utils import is_hpu
 
 logger = init_logger(__name__)
-if is_hpu():
- try:
- from habana_frameworks.torch.hpex.normalization import (FusedRMSNorm as
- HPUFusedRMSNorm
- )
- except ImportError:
- logger.warning(
- "Could not import HPU FusedRMSNorm kernel. "
- "vLLM will use forward_native implementation of RMSNorm.")
- HPUFusedRMSNorm = None
 
 
 class RMSNorm(CustomOp):
@@ -86,6 +75,7 @@ def forward_hpu(
  x: torch.Tensor,
  residual: Optional[torch.Tensor] = None,
  ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ from vllm.hpu.ops import HPUFusedRMSNorm
  if HPUFusedRMSNorm is None:
  return self.forward_native(x, residual)
  if residual is not None:

vllm/utils.py

@@ -664,7 +664,7 @@ def get_summary_string(self):
  return (
  f"{format_bytes(self.consumed_device_memory)} of device memory "
  f"({format_bytes(self.final_device_memory)}/"
- f"({format_bytes(HabanaMemoryProfiler.total_device_memory())} used)"
+ f"{format_bytes(HabanaMemoryProfiler.total_device_memory())} used)"
  f" and {format_bytes(self.consumed_host_memory)} of host memory "
  f"({format_bytes(self.final_host_memory)}/"
  f"{format_bytes(HabanaMemoryProfiler.total_host_memory())} used)")

vllm/worker/habana_model_runner.py

@@ -409,7 +409,7 @@ def __init__(
 
  # Profiler stats
  self.profiler_counter_helper = HabanaProfilerCounterHelper()
-
+ self._mem_margin: Optional[int] = None
  self._setup_buckets()
 
  def load_model(self) -> None:
@@ -1071,10 +1071,15 @@ def warmup_all_buckets(self, buckets, is_prompt, kv_caches):
  len(buckets), batch_size, seq_len)
  self.warmup_scenario(batch_size, seq_len, is_prompt, kv_caches)
 
- def warmup_graphs(self, strategy, buckets, is_prompt, kv_caches,
- available_mem):
- total_batch_seq = 0.001
- total_mem = 0
+ def warmup_graphs(self,
+ strategy,
+ buckets,
+ is_prompt,
+ kv_caches,
+ available_mem,
+ starting_mem=0,
+ total_batch_seq=0.001):
+ total_mem = starting_mem
  idx = 0
  phase = f'Graph/{"Prompt" if is_prompt else "Decode"}'
  num_candidates = len(buckets)
@@ -1088,14 +1093,18 @@ def warmup_graphs(self, strategy, buckets, is_prompt, kv_caches,
  raise NotImplementedError(
  f'Unsupported graph allocation strategy: {strategy}')
  buckets = list(sorted(buckets, key=ordering))
-
+ captured_all = True
  for idx, (batch_size, seq_len) in enumerate(buckets):
  # Graph memory usage is proportional to seq dimension in a batch
  batch_seq = batch_size * seq_len if is_prompt else batch_size
  mem_estimate = batch_seq / total_batch_seq * total_mem
  if mem_estimate >= available_mem:
+ captured_all = False
+ continue
+ graphed_bucket = (batch_size, seq_len, is_prompt)
+ if graphed_bucket in self.graphed_buckets:
  continue
- self.graphed_buckets.add((batch_size, seq_len, is_prompt))
+ self.graphed_buckets.add(graphed_bucket)
  self.log_warmup(phase, idx, num_candidates, batch_size, seq_len)
  with HabanaMemoryProfiler() as mem_prof:
  self.warmup_scenario(batch_size, seq_len, is_prompt, kv_caches)
@@ -1104,6 +1113,12 @@ def warmup_graphs(self, strategy, buckets, is_prompt, kv_caches,
  available_mem -= used_mem
  total_mem += used_mem
  total_batch_seq += batch_seq
+
+ return total_mem, total_batch_seq, captured_all
+
+ def log_graph_warmup_summary(self, buckets, is_prompt, total_mem):
+ num_candidates = len(buckets)
+ phase = f'Graph/{"Prompt" if is_prompt else "Decode"}'
  graphed = list(c[:2] for c in self.graphed_buckets
  if c[2] == is_prompt)
  msg = (f'{phase} captured:{len(graphed)} '
@@ -1124,22 +1139,63 @@ def warmup_model(self, kv_caches: List[torch.Tensor]) -> None:
  self.warmup_all_buckets(self.decode_buckets, False, kv_caches)
 
  if not self.enforce_eager and htorch.utils.internal.is_lazy():
- mem_margin = 1.0 - float(
- os.environ.get('VLLM_GRAPH_MEM_MARGIN', '0.02'))
- free_mem = \
- mem_margin * HabanaMemoryProfiler.current_free_device_memory()
- free_mem = align_workers(free_mem, torch.distributed.ReduceOp.MIN)
+ assert self.mem_margin is not None, \
+ ("HabanaWorker.determine_num_available_blocks needs "
+ "to be called before warming up the model.")
+ free_mem = HabanaMemoryProfiler.current_free_device_memory()
+ graph_free_mem = free_mem - self.mem_margin
+ graph_free_mem = align_workers(graph_free_mem,
+ torch.distributed.ReduceOp.MIN)
  prompt_graph_mem_ratio = float(
  os.environ.get('VLLM_GRAPH_PROMPT_RATIO', '0.5'))
- prompt_available_memory = prompt_graph_mem_ratio * free_mem
- decode_available_memory = free_mem - prompt_available_memory
- prompt_strategy = 'min_tokens'
+ prompt_available_memory = prompt_graph_mem_ratio * graph_free_mem
+ decode_available_memory = graph_free_mem - prompt_available_memory
+ msg = (f"Using {format_bytes(graph_free_mem)}"
+ f"/{format_bytes(free_mem)} "
+ "of free device memory for HPUGraphs, "
+ f"{format_bytes(prompt_available_memory)} for prompt and "
+ f"{format_bytes(decode_available_memory)} for decode "
+ f"(VLLM_GRAPH_PROMPT_RATIO={prompt_graph_mem_ratio})")
+ logger.info(msg)
+ prompt_strategy = os.environ.get('VLLM_GRAPH_PROMPT_STRATEGY',
+ 'min_tokens')
  decode_strategy = os.environ.get('VLLM_GRAPH_DECODE_STRATEGY',
  'max_bs')
- self.warmup_graphs(prompt_strategy, self.prompt_buckets, True,
- kv_caches, prompt_available_memory)
- self.warmup_graphs(decode_strategy, self.decode_buckets, False,
- kv_caches, decode_available_memory)
+ mem_post_prompt, prompt_batch_seq, prompt_captured_all = \
+ self.warmup_graphs(
+ prompt_strategy, self.prompt_buckets, True, kv_caches,
+ prompt_available_memory)
+ mem_post_decode, decode_batch_seq, decode_captured_all = \
+ self.warmup_graphs(
+ decode_strategy, self.decode_buckets, False, kv_caches,
+ decode_available_memory)
+
+ # Not all prompt buckets were captured, but all decode buckets were
+ # captured and we have some free graph-allocated space left.
+ # Let's try to use it for capturing more prompt buckets.
+ if mem_post_decode + mem_post_prompt < graph_free_mem \
+ and not prompt_captured_all \
+ and decode_captured_all:
+ mem_post_prompt, _, prompt_captured_all = self.warmup_graphs(
+ prompt_strategy, self.prompt_buckets, True, kv_caches,
+ graph_free_mem - mem_post_prompt - mem_post_decode,
+ mem_post_prompt, prompt_batch_seq)
+
+ # Not all decode buckets were captured, but all prompt buckets were
+ # captured and we have some free graph-allocated space left.
+ # Let's try to use it for capturing more decode buckets.
+ if mem_post_decode + mem_post_prompt < graph_free_mem \
+ and not decode_captured_all \
+ and prompt_captured_all:
+ mem_post_decode, _, _ = self.warmup_graphs(
+ decode_strategy, self.decode_buckets, False, kv_caches,
+ graph_free_mem - mem_post_prompt - mem_post_decode,
+ mem_post_decode, decode_batch_seq)
+
+ self.log_graph_warmup_summary(self.prompt_buckets, True,
+ mem_post_prompt)
+ self.log_graph_warmup_summary(self.decode_buckets, False,
+ mem_post_decode)
 
  end_time = time.perf_counter()
  end_mem = HabanaMemoryProfiler.current_device_memory_usage()
@@ -1154,6 +1210,14 @@ def warmup_model(self, kv_caches: List[torch.Tensor]) -> None:
  def vocab_size(self) -> int:
  return self.model_config.get_vocab_size()
 
+ @property
+ def mem_margin(self) -> Optional[int]:
+ return self._mem_margin
+
+ @mem_margin.setter
+ def mem_margin(self, value):
+ self._mem_margin = value
+
 
 def _maybe_wrap_in_hpu_graph(*args, **kwargs):
  return htorch.hpu.wrap_in_hpu_graph(HpuModelAdapter(