Add beam_search ability to sharktank LLM model

nod-ai · Feb 21, 2025 · c28bf2a · c28bf2a
1 parent 888a98a
commit c28bf2a
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 4 deletions.
diff --git a/sharktank/sharktank/examples/export_paged_llm_v1.py b/sharktank/sharktank/examples/export_paged_llm_v1.py
@@ -83,6 +83,7 @@ def main():
         block_seq_stride=args.block_seq_stride,
         activation_dtype=args.activation_dtype,
         attention_dtype=args.attention_dtype,
+        n_beams=args.n_beams,
     )
     llama_config.fake_quant = args.fake_quant
 
@@ -112,6 +113,7 @@ def generate_params_json(
             "prefill_batch_sizes": prefill_bs,
             "decode_batch_sizes": decode_bs,
             "transformer_block_count": hp.block_count,
+            "n_beams": llama_config.n_beams,
             "paged_kv_cache": {
                 "attention_head_count_kv": hp.attention_head_count_kv,
                 "block_seq_stride": llama_config.block_seq_stride,
@@ -251,14 +253,15 @@ def generate_batch_decode(bs: int):
         block_dim_min = 2
         block_dim_max = ceildiv(hp.context_length, llama_config.block_seq_stride) - 1
         block_dim = torch.export.Dim("block", min=block_dim_min, max=block_dim_max)
+        batch_dim_size = bs * llama_config.n_beams
         tokens = torch.empty(
-            bs,
+            batch_dim_size,
             1,
             dtype=torch.int64,
         )
-        seq_lens = torch.empty(bs, dtype=torch.int64)
-        start_positions = torch.ones(bs, dtype=torch.int64)
-        seq_block_ids = torch.empty(bs, block_dim_min, dtype=torch.int64)
+        seq_lens = torch.empty(batch_dim_size, dtype=torch.int64)
+        start_positions = torch.ones(batch_dim_size, dtype=torch.int64)
+        seq_block_ids = torch.empty(batch_dim_size, block_dim_min, dtype=torch.int64)
 
         (
             cache_state,

diff --git a/sharktank/sharktank/layers/configs/llm_configs.py b/sharktank/sharktank/layers/configs/llm_configs.py
@@ -191,6 +191,11 @@ class LlamaModelConfig:
     # the program and not.
     static_tables: bool = True
 
+    # The number of beams to use when generating tokens for a given prompt.
+    # When n_beams == 1, `greedy` selection is used,
+    # when n_beams > 1, `beam search` is used.
+    n_beams: int = 1
+
 
 @dataclass
 class T5Config:

diff --git a/sharktank/sharktank/utils/cli.py b/sharktank/sharktank/utils/cli.py
@@ -120,6 +120,12 @@ def add_model_options(parser: argparse.ArgumentParser):
         type=int,
         default=512,
     )
+    parser.add_argument(
+        "--n-beams",
+        help="Number of beams to use when generating tokens.",
+        type=int,
+        default=1,
+    )
 
 
 def add_quantization_options(parser: argparse.ArgumentParser):