We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Investigate top-p performance
Currently our top-p (nucleus) sampler has the worst performance of our entire sampling offering, likely because of the sort operation we need to do for our softmax probabilities. See [these rough numbers](https://github.com/keras-team/keras-nlp/commit/41af01200c5d9898f381477b6a67a87b6951f844#diff-576c27794bae284efc278336a16d0649aaf599973a64d4b65dc67bc99bb7f000) for example, where top-p is slower than beam, despite needing much less computation on the model itself.

We should see if there is anything we can do to speed up top-p sampling with XLA (e.g. a more XLA friendly sort op).

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `keras_nlp/samplers/top_p_sampler.py`
Content:
```
1 # Copyright 2023 The KerasNLP Authors
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     https://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 """Top-p Sampler."""
15 
16 import tensorflow as tf
17 
18 from keras_nlp.api_export import keras_nlp_export
19 from keras_nlp.samplers.sampler import Sampler
20 from keras_nlp.samplers.sampler import call_args_docstring
21 from keras_nlp.utils.python_utils import format_docstring
22 
23 
24 @format_docstring(call_args=call_args_docstring)
25 @keras_nlp_export("keras_nlp.samplers.TopPSampler")
26 class TopPSampler(Sampler):
27     """Top-P Sampler class.
28 
29     This sampler implements top-p search algorithm. Top-p search selects tokens
30     from the smallest subset of output probabilities that sum to greater than
31     `p`. Put in another way, top-p will first order token predictions by
32     likelihood, and ignore all tokens after the cumulative probability of
33     selected tokens exceeds `p`, then select a token from the remaining tokens.
34 
35     Args:
36         p: float, the `p` value of top-p.
37         seed: int, defaults to None. The random seed.
38 
39     Call Args:
40         {{call_args}}
41 
42     Examples:
43     ```python
44     # Use a simple alphabet of lowercase characters with ids in range [0, 25].
45     int_lookup = {i: chr(i + ord('a')) for i in range(26)}
46     char_lookup = {v: k for k, v in int_lookup.items()}
47     batch_size, length, vocab_size = 1, 12, len(int_lookup)
48 
49     def next(prompt, cache, index):
50         hidden_states = tf.ones((batch_size, 10))
51         # A uniform distribution over our alphabet.
52         logits = tf.ones((batch_size, vocab_size))
53         return logits, hidden_states, cache
54 
55     output = keras_nlp.samplers.TopPSampler(p=0.1)(
56         next=next,
57         prompt=tf.fill((batch_size, length,), char_lookup['z']),
58         index=5,
59     )
60     print(["".join([int_lookup[i] for i in s]) for s in output.numpy()])
61     # >>> ['zzzzzbabcccb']
62     ```
63     """
64 
65     def __init__(
66         self,
67         p=0.1,
68         seed=None,
69     ):
70         super().__init__()
71         self.p = p
72         self.seed = seed
73 
74     def get_next_token(self, probabilities):
75         # Sort preds in descending order.
76         sorted_preds, sorted_indices = tf.math.top_k(
77             probabilities, k=tf.shape(probabilities)[1], sorted=True
78         )
79         # Calculate cumulative probability distribution.
80         cumulative_probabilities = tf.math.cumsum(sorted_preds, axis=-1)
81         # Create a mask for the tokens to keep.
82         keep_mask = cumulative_probabilities <= self.p
83         # Shift to include the last token that exceed p.
84         shifted_keep_mask = tf.concat(
85             [tf.ones_like(keep_mask[:, :1]), keep_mask[:, :-1]], axis=-1
86         )
87         # Filter out unmasked tokens and sample from filtered distribution.
88         probabilities = tf.where(
89             shifted_keep_mask,
90             sorted_preds,
91             tf.zeros(tf.shape(probabilities), dtype=sorted_preds.dtype),
92         )
93         sorted_next_token = tf.random.categorical(
94             tf.math.log(probabilities), 1, seed=self.seed
95         )
96         return tf.gather_nd(sorted_indices, sorted_next_token, batch_dims=1)
97 
98     def get_config(self):
99         config = super().get_config()
100         config.update(
101             {
102                 "p": self.p,
103                 "seed": self.seed,
104             }
105         )
106         return config
107
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/keras_nlp/samplers/top_p_sampler.py b/keras_nlp/samplers/top_p_sampler.py
--- a/keras_nlp/samplers/top_p_sampler.py
+++ b/keras_nlp/samplers/top_p_sampler.py
@@ -34,6 +34,12 @@
 
     Args:
         p: float, the `p` value of top-p.
+        k: int, defaults to None. If set, this argument defines a
+            heuristic "top-k" cutoff applied before the "top-p" sampling. All
+            logits not in the top `k` will be discarded, and the remaining
+            logits will be sorted to find a cutoff point for `p`. Setting this
+            arg can significantly speed sampling up by reducing the number
+            of tokens to sort.
         seed: int, defaults to None. The random seed.
 
     Call Args:
@@ -65,16 +71,21 @@
     def __init__(
         self,
         p=0.1,
+        k=None,
         seed=None,
     ):
         super().__init__()
         self.p = p
+        self.k = k
         self.seed = seed
 
     def get_next_token(self, probabilities):
-        # Sort preds in descending order.
+        cutoff = tf.shape(probabilities)[1]
+        if self.k is not None:
+            # If `k` is set, only sample from top `k` tokens.
+            cutoff = tf.math.minimum(cutoff, self.k)
         sorted_preds, sorted_indices = tf.math.top_k(
-            probabilities, k=tf.shape(probabilities)[1], sorted=True
+            probabilities, k=cutoff, sorted=True
         )
         # Calculate cumulative probability distribution.
         cumulative_probabilities = tf.math.cumsum(sorted_preds, axis=-1)
@@ -88,7 +99,7 @@
         probabilities = tf.where(
             shifted_keep_mask,
             sorted_preds,
-            tf.zeros(tf.shape(probabilities), dtype=sorted_preds.dtype),
+            tf.zeros(tf.shape(sorted_preds), dtype=sorted_preds.dtype),
         )
         sorted_next_token = tf.random.categorical(
             tf.math.log(probabilities), 1, seed=self.seed
@@ -100,6 +111,7 @@
         config.update(
             {
                 "p": self.p,
+                "k": self.k,
                 "seed": self.seed,
             }
         )

{"golden_diff": "diff --git a/keras_nlp/samplers/top_p_sampler.py b/keras_nlp/samplers/top_p_sampler.py\n--- a/keras_nlp/samplers/top_p_sampler.py\n+++ b/keras_nlp/samplers/top_p_sampler.py\n@@ -34,6 +34,12 @@\n \n     Args:\n         p: float, the `p` value of top-p.\n+        k: int, defaults to None. If set, this argument defines a\n+            heuristic \"top-k\" cutoff applied before the \"top-p\" sampling. All\n+            logits not in the top `k` will be discarded, and the remaining\n+            logits will be sorted to find a cutoff point for `p`. Setting this\n+            arg can significantly speed sampling up by reducing the number\n+            of tokens to sort.\n         seed: int, defaults to None. The random seed.\n \n     Call Args:\n@@ -65,16 +71,21 @@\n     def __init__(\n         self,\n         p=0.1,\n+        k=None,\n         seed=None,\n     ):\n         super().__init__()\n         self.p = p\n+        self.k = k\n         self.seed = seed\n \n     def get_next_token(self, probabilities):\n-        # Sort preds in descending order.\n+        cutoff = tf.shape(probabilities)[1]\n+        if self.k is not None:\n+            # If `k` is set, only sample from top `k` tokens.\n+            cutoff = tf.math.minimum(cutoff, self.k)\n         sorted_preds, sorted_indices = tf.math.top_k(\n-            probabilities, k=tf.shape(probabilities)[1], sorted=True\n+            probabilities, k=cutoff, sorted=True\n         )\n         # Calculate cumulative probability distribution.\n         cumulative_probabilities = tf.math.cumsum(sorted_preds, axis=-1)\n@@ -88,7 +99,7 @@\n         probabilities = tf.where(\n             shifted_keep_mask,\n             sorted_preds,\n-            tf.zeros(tf.shape(probabilities), dtype=sorted_preds.dtype),\n+            tf.zeros(tf.shape(sorted_preds), dtype=sorted_preds.dtype),\n         )\n         sorted_next_token = tf.random.categorical(\n             tf.math.log(probabilities), 1, seed=self.seed\n@@ -100,6 +111,7 @@\n         config.update(\n             {\n                 \"p\": self.p,\n+                \"k\": self.k,\n                 \"seed\": self.seed,\n             }\n         )\n", "issue": "Investigate top-p performance\nCurrently our top-p (nucleus) sampler has the worst performance of our entire sampling offering, likely because of the sort operation we need to do for our softmax probabilities. See [these rough numbers](https://github.com/keras-team/keras-nlp/commit/41af01200c5d9898f381477b6a67a87b6951f844#diff-576c27794bae284efc278336a16d0649aaf599973a64d4b65dc67bc99bb7f000) for example, where top-p is slower than beam, despite needing much less computation on the model itself.\r\n\r\nWe should see if there is anything we can do to speed up top-p sampling with XLA (e.g. a more XLA friendly sort op).\n", "before_files": [{"content": "# Copyright 2023 The KerasNLP Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Top-p Sampler.\"\"\"\n\nimport tensorflow as tf\n\nfrom keras_nlp.api_export import keras_nlp_export\nfrom keras_nlp.samplers.sampler import Sampler\nfrom keras_nlp.samplers.sampler import call_args_docstring\nfrom keras_nlp.utils.python_utils import format_docstring\n\n\n@format_docstring(call_args=call_args_docstring)\n@keras_nlp_export(\"keras_nlp.samplers.TopPSampler\")\nclass TopPSampler(Sampler):\n    \"\"\"Top-P Sampler class.\n\n    This sampler implements top-p search algorithm. Top-p search selects tokens\n    from the smallest subset of output probabilities that sum to greater than\n    `p`. Put in another way, top-p will first order token predictions by\n    likelihood, and ignore all tokens after the cumulative probability of\n    selected tokens exceeds `p`, then select a token from the remaining tokens.\n\n    Args:\n        p: float, the `p` value of top-p.\n        seed: int, defaults to None. The random seed.\n\n    Call Args:\n        {{call_args}}\n\n    Examples:\n    ```python\n    # Use a simple alphabet of lowercase characters with ids in range [0, 25].\n    int_lookup = {i: chr(i + ord('a')) for i in range(26)}\n    char_lookup = {v: k for k, v in int_lookup.items()}\n    batch_size, length, vocab_size = 1, 12, len(int_lookup)\n\n    def next(prompt, cache, index):\n        hidden_states = tf.ones((batch_size, 10))\n        # A uniform distribution over our alphabet.\n        logits = tf.ones((batch_size, vocab_size))\n        return logits, hidden_states, cache\n\n    output = keras_nlp.samplers.TopPSampler(p=0.1)(\n        next=next,\n        prompt=tf.fill((batch_size, length,), char_lookup['z']),\n        index=5,\n    )\n    print([\"\".join([int_lookup[i] for i in s]) for s in output.numpy()])\n    # >>> ['zzzzzbabcccb']\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        p=0.1,\n        seed=None,\n    ):\n        super().__init__()\n        self.p = p\n        self.seed = seed\n\n    def get_next_token(self, probabilities):\n        # Sort preds in descending order.\n        sorted_preds, sorted_indices = tf.math.top_k(\n            probabilities, k=tf.shape(probabilities)[1], sorted=True\n        )\n        # Calculate cumulative probability distribution.\n        cumulative_probabilities = tf.math.cumsum(sorted_preds, axis=-1)\n        # Create a mask for the tokens to keep.\n        keep_mask = cumulative_probabilities <= self.p\n        # Shift to include the last token that exceed p.\n        shifted_keep_mask = tf.concat(\n            [tf.ones_like(keep_mask[:, :1]), keep_mask[:, :-1]], axis=-1\n        )\n        # Filter out unmasked tokens and sample from filtered distribution.\n        probabilities = tf.where(\n            shifted_keep_mask,\n            sorted_preds,\n            tf.zeros(tf.shape(probabilities), dtype=sorted_preds.dtype),\n        )\n        sorted_next_token = tf.random.categorical(\n            tf.math.log(probabilities), 1, seed=self.seed\n        )\n        return tf.gather_nd(sorted_indices, sorted_next_token, batch_dims=1)\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"p\": self.p,\n                \"seed\": self.seed,\n            }\n        )\n        return config\n", "path": "keras_nlp/samplers/top_p_sampler.py"}], "after_files": [{"content": "# Copyright 2023 The KerasNLP Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Top-p Sampler.\"\"\"\n\nimport tensorflow as tf\n\nfrom keras_nlp.api_export import keras_nlp_export\nfrom keras_nlp.samplers.sampler import Sampler\nfrom keras_nlp.samplers.sampler import call_args_docstring\nfrom keras_nlp.utils.python_utils import format_docstring\n\n\n@format_docstring(call_args=call_args_docstring)\n@keras_nlp_export(\"keras_nlp.samplers.TopPSampler\")\nclass TopPSampler(Sampler):\n    \"\"\"Top-P Sampler class.\n\n    This sampler implements top-p search algorithm. Top-p search selects tokens\n    from the smallest subset of output probabilities that sum to greater than\n    `p`. Put in another way, top-p will first order token predictions by\n    likelihood, and ignore all tokens after the cumulative probability of\n    selected tokens exceeds `p`, then select a token from the remaining tokens.\n\n    Args:\n        p: float, the `p` value of top-p.\n        k: int, defaults to None. If set, this argument defines a\n            heuristic \"top-k\" cutoff applied before the \"top-p\" sampling. All\n            logits not in the top `k` will be discarded, and the remaining\n            logits will be sorted to find a cutoff point for `p`. Setting this\n            arg can significantly speed sampling up by reducing the number\n            of tokens to sort.\n        seed: int, defaults to None. The random seed.\n\n    Call Args:\n        {{call_args}}\n\n    Examples:\n    ```python\n    # Use a simple alphabet of lowercase characters with ids in range [0, 25].\n    int_lookup = {i: chr(i + ord('a')) for i in range(26)}\n    char_lookup = {v: k for k, v in int_lookup.items()}\n    batch_size, length, vocab_size = 1, 12, len(int_lookup)\n\n    def next(prompt, cache, index):\n        hidden_states = tf.ones((batch_size, 10))\n        # A uniform distribution over our alphabet.\n        logits = tf.ones((batch_size, vocab_size))\n        return logits, hidden_states, cache\n\n    output = keras_nlp.samplers.TopPSampler(p=0.1)(\n        next=next,\n        prompt=tf.fill((batch_size, length,), char_lookup['z']),\n        index=5,\n    )\n    print([\"\".join([int_lookup[i] for i in s]) for s in output.numpy()])\n    # >>> ['zzzzzbabcccb']\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        p=0.1,\n        k=None,\n        seed=None,\n    ):\n        super().__init__()\n        self.p = p\n        self.k = k\n        self.seed = seed\n\n    def get_next_token(self, probabilities):\n        cutoff = tf.shape(probabilities)[1]\n        if self.k is not None:\n            # If `k` is set, only sample from top `k` tokens.\n            cutoff = tf.math.minimum(cutoff, self.k)\n        sorted_preds, sorted_indices = tf.math.top_k(\n            probabilities, k=cutoff, sorted=True\n        )\n        # Calculate cumulative probability distribution.\n        cumulative_probabilities = tf.math.cumsum(sorted_preds, axis=-1)\n        # Create a mask for the tokens to keep.\n        keep_mask = cumulative_probabilities <= self.p\n        # Shift to include the last token that exceed p.\n        shifted_keep_mask = tf.concat(\n            [tf.ones_like(keep_mask[:, :1]), keep_mask[:, :-1]], axis=-1\n        )\n        # Filter out unmasked tokens and sample from filtered distribution.\n        probabilities = tf.where(\n            shifted_keep_mask,\n            sorted_preds,\n            tf.zeros(tf.shape(sorted_preds), dtype=sorted_preds.dtype),\n        )\n        sorted_next_token = tf.random.categorical(\n            tf.math.log(probabilities), 1, seed=self.seed\n        )\n        return tf.gather_nd(sorted_indices, sorted_next_token, batch_dims=1)\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"p\": self.p,\n                \"k\": self.k,\n                \"seed\": self.seed,\n            }\n        )\n        return config\n", "path": "keras_nlp/samplers/top_p_sampler.py"}]}