fix and add test

integration-tests/test_classifications.py

@@ -16,7 +16,33 @@ def test_distilbert_ner():
     with run_lorax_container(config):
         response = requests.post(
             "http://localhost:8080/classify",
-            json={"inputs": "Johnny supports the Golden State Warriors. He lives in London."},
+            json={
+                "inputs": "Johnny supports the Golden State Warriors. He lives in London."
+            },
+        )
+        response.raise_for_status()
+        print("RESPONSE FROM CLASSIFICATION:", response.json())
+        assert len(response.json()) > 0
+
+
+def test_bert_ner():
+    config = {
+        "name": "bert-ner",
+        "model_id": "dslim/bert-base-NER",
+        "docker_args": {
+            "max_input_length": 512,
+            "max_batch_prefill_tokens": 512,
+            "max_batch_total_tokens": 512,
+            "max_total_tokens": 512,
+            "backend": "flashinfer",
+        },
+    }
+    with run_lorax_container(config):
+        response = requests.post(
+            "http://localhost:8080/classify",
+            json={
+                "inputs": "Johnny supports the Golden State Warriors. He lives in London."
+            },
         )
         response.raise_for_status()
         print("RESPONSE FROM CLASSIFICATION:", response.json())

server/lorax_server/models/flash_bert.py

@@ -179,6 +179,8 @@ def _forward_context(
             num_heads=self.num_heads,
             num_kv_heads=self.num_kv_heads,
             head_size=self.head_size,
+            # TODO: This is a hack to get the prefill state to work
+            window_left=0,
         )
 
     def forward(self, batch: FlashEmbeddingClassificationBatch):