Reinstate llama tests (#141)

Author: alessandropalla

intel_npu_acceleration_library/_version.py

@@ -3,4 +3,4 @@
 # SPDX-License-Identifier: Apache 2.0
 #
 
-__version__ = "v1.3.0"
+__version__ = "v1.4.0"

intel_npu_acceleration_library/nn/llm.py

@@ -154,7 +154,7 @@ def __init__(
 
         self.hidden_size = config.hidden_size
         self.num_heads = config.num_attention_heads
-        self.head_dim = self.hidden_size // self.num_heads
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
         self.num_key_value_heads = config.num_key_value_heads
         self.num_key_value_groups = self.num_heads // self.num_key_value_heads
         self.is_causal = True
@@ -238,7 +238,7 @@ def forward(
         )
 
         attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = attn_output.view(bsz, q_len, -1)
 
         attn_output = self.o_proj(attn_output)
 

test/python/test_optimizations.py

@@ -81,6 +81,7 @@ def get_model(model_name, hidden_size, intermediate_size, bias):
         conf.num_hidden_layers = 1
         conf.hidden_size = hidden_size
         conf.intermediate_size = intermediate_size
+        conf.head_dim = conf.hidden_size // conf.num_attention_heads
 
         return LlamaModel(conf)
     elif model_name == "GemmaModel":
@@ -130,9 +131,6 @@ def test_fusion(model_name, hidden_size, intermediate_size, batch, bias):
 @pytest.mark.parametrize("bias", [True, False])
 def test_model(model_name, hidden_size, intermediate_size, sequence_length, bias):
 
-    if model_name == "LlamaModel":
-        pytest.skip("LlamaModel Fix in progress")
-
     with torch.no_grad():
         model = get_model(model_name, hidden_size, intermediate_size, bias).eval()
         example_input = torch.randint(