nice code

Signed-off-by: Liu, Kaixuan <[email protected]>

Author: kaixuanliu

backends/python/server/text_embeddings_server/models/flash_bert.py

@@ -12,7 +12,7 @@
 from text_embeddings_server.models import Model
 from text_embeddings_server.models.types import FlashBatch, Embedding
 from text_embeddings_server.utils.flash_attn import attention
-
+from text_embeddings_server.utils.device import use_ipex
 tracer = trace.get_tracer(__name__)
 
 
@@ -25,6 +25,8 @@ def __init__(self, prefix, handle, device, dtype, config: BertConfig):
 
     def forward(self, hidden_states, residual=None):
         # Flash attention imports
+        normed_hidden_states = None
+        res = None
         if self.device.type == "cuda":
             import dropout_layer_norm
             normed_hidden_states, res, *rest = dropout_layer_norm.dropout_add_ln_fwd(
@@ -46,7 +48,7 @@ def forward(self, hidden_states, residual=None):
             )
             if res is None:
                 res = hidden_states
-        else:
+        elif use_ipex():
             import intel_extension_for_pytorch as ipex
             normed_hidden_states = ipex.llm.functional.add_layer_norm(
                 residual,

backends/python/server/text_embeddings_server/utils/device.py

@@ -32,16 +32,7 @@ def _is_hpu() -> bool:
     try:
         subprocess.run(["hl-smi"], capture_output=True, check=True)
     except (FileNotFoundError, PermissionError, subprocess.CalledProcessError):
-        if not os.path.exists('/dev/accel/accel0') and not os.path.exists(
-                '/dev/accel/accel_controlD0'):
-            # last resort...
-            try:
-                output = subprocess.check_output(
-                    'lsmod | grep habanalabs | wc -l', shell=True)
-                is_hpu_available = int(output) > 0
-            except (ValueError, FileNotFoundError, PermissionError,
-                    subprocess.CalledProcessError):
-                is_hpu_available = False
+        is_hpu_available = False
     return is_hpu_available
 
 def use_ipex() -> bool:

backends/python/server/text_embeddings_server/utils/flash_attn.py

@@ -58,7 +58,10 @@ def attention(q, k, v, out, cu_seqlens, max_s, softmax_scale, is_causal=False):
     if HAS_FLASH_ATTN_V2:
         if use_ipex():
             import intel_extension_for_pytorch as ipex
-            return ipex.llm.functional.varlen_attention(q, k, v, out, cu_seqlens, cu_seqlens, max_s, max_s, 0, softmax_scale, zero_tensors=False, is_causal=False, return_softmax=False, gen_=None)
+            return ipex.llm.functional.varlen_attention(q, k, v, out, cu_seqlens, cu_seqlens,
+                                                        max_s, max_s, 0, softmax_scale,
+                                                        zero_tensors=False, is_causal=False,
+                                                        return_softmax=False, gen_=None)
         else:
             return flash_attn_2_cuda.varlen_fwd(
                 q,