Llama export TP example with KV cache

tools/llm/tensor_parallel_llama_export.py

@@ -1,5 +1,5 @@
 """
-Tensor Parallel Llama: torch.export → save → load TRT engines across nodes.
+Tensor Parallel LLM (Llama / Qwen): torch.export → save → load TRT engines.
 
 Two modes:
   export  — torch.export → TRT AOT compile → save per-rank engines to disk
@@ -11,32 +11,53 @@
 the weights per-rank, injecting NCCL all-reduce ops via Torch-TensorRT's
 distributed compilation path.
 
+The manual slicing assumes the standard decoder layout
+(``model.model.layers[i].self_attn.{q,k,v,o}_proj`` +
+``mlp.{gate,up,down}_proj``), which both Llama and Qwen2/Qwen3 follow, so the
+same script works for either.
 
+Tested ungated models (no HF token needed):
+  Qwen/Qwen2.5-0.5B-Instruct  - smallest/fastest smoke test; kv_heads=2 => TP=2 ONLY
+  Qwen/Qwen3-1.7B             - default; kv_heads=8 => TP 2/4/8
+  meta-llama/Llama-3.2-1B-Instruct - original example (GATED, needs HF token)
 
 Usage
 -----
-# Export mode — export + compile + save engines (run on each node):
-  torchtrtrun --nproc_per_node=1 --nnodes=2 --node_rank=0 \\
-      --rdzv_endpoint=<node0-ip>:29500 \\
-      tools/llm/tensor_parallel_llama_export.py --mode export --save_dir /tmp/llama_tp_engines
+# Export mode — single node, 2 GPUs (ungated Qwen, no token needed):
+  torchtrtrun --nproc_per_node=2 \\
+      tools/llm/tensor_parallel_llama_export.py \\
+      --mode export --save_dir /tmp/llm_tp_engines --model Qwen/Qwen3-1.7B
 
 # Load mode — load saved engines + inference (no model download needed):
+  torchtrtrun --nproc_per_node=2 \\
+      tools/llm/tensor_parallel_llama_export.py \\
+      --mode load --save_dir /tmp/llm_tp_engines --model Qwen/Qwen3-1.7B
+
+# Multi-node export (one GPU per node):
   torchtrtrun --nproc_per_node=1 --nnodes=2 --node_rank=0 \\
       --rdzv_endpoint=<node0-ip>:29500 \\
-      tools/llm/tensor_parallel_llama_export.py --mode load --save_dir /tmp/llama_tp_engines
+      tools/llm/tensor_parallel_llama_export.py --mode export --save_dir /tmp/llm_tp_engines
 """
 
 import argparse
 import datetime
 import logging
 import os
 import sys
+import timeit
+from contextlib import ExitStack
 from pathlib import Path
 
 import torch
 import torch.distributed as dist
 import torch.distributed.tensor._dtensor_spec
 import torch.utils._pytree
+from utils import (
+    generate,
+    generate_with_static_cache,
+    record_stats_split,
+    time_generate_split,
+)
 
 # DTensorSpec must be a pytree constant before torch.export traces a TP model.
 torch.utils._pytree.register_constant(
@@ -59,13 +80,17 @@
 
 setup_nccl_for_torch_tensorrt()
 
+from torchtrt_ext import register_sdpa
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 logging.basicConfig(
     level=logging.INFO,
     format=f"[Rank {rank}] %(levelname)s: %(message)s",
 )
 logger = logging.getLogger(__name__)
+# Quiet torch_tensorrt's default verbose output. `--debug` still re-enables
+# Debug-level logging during compile via torch_tensorrt.logging.debug().
+torch_tensorrt.logging.set_level(logging.ERROR)
 logger.info(f"dist init OK  rank={rank}/{world_size}  device={DEVICE}")
 
 
@@ -87,18 +112,39 @@ def _extract_logits(outputs):
     return outputs
 
 
-def generate_greedy(model, input_ids, max_len, eos_token_id):
-    """Greedy decode that works with any model output format."""
-    seq = input_ids.clone()
-    for _ in range(max_len - input_ids.shape[1]):
-        position_ids = torch.arange(seq.shape[1]).unsqueeze(0).to(seq.device)
-        outputs = model(seq, position_ids=position_ids)
-        logits = _extract_logits(outputs)
-        next_token = logits[:, -1, :].argmax(dim=-1)
-        seq = torch.cat([seq, next_token[:, None]], dim=-1)
-        if (next_token == eos_token_id).all():
-            break
-    return seq
+def time_generate(generate_fn, model, input_ids, max_len, eos_token_id, iterations=5):
+    """Measure end-to-end generation latency over multiple iterations."""
+    timings = []
+    for _ in range(iterations):
+        start = timeit.default_timer()
+        generate_fn(model, input_ids.clone(), max_len, eos_token_id)
+        torch.cuda.synchronize()
+        timings.append(timeit.default_timer() - start)
+    return timings
+
+
+def _pick_generate_fn(args):
+    """Pick the right greedy-decode helper based on the --cache flag."""
+    if args.cache in ("static_v1", "static_v2"):
+        return generate_with_static_cache
+    return generate
+
+
+def record_stats(backend, timings, precision, batch_size=1):
+    import numpy as np
+
+    times = np.array(timings)
+    speeds = batch_size / times
+    return {
+        "Backend": backend,
+        "Model Precision": precision,
+        "Batch size": batch_size,
+        "Median(FPS)": float(np.median(speeds)),
+        "Mean(FPS)": float(np.mean(speeds)),
+        "Median-Latency(ms)": float(np.median(times)) * 1000,
+        "Mean-Latency(ms)": float(np.mean(times)) * 1000,
+        "Latency-StdDev(ms)": float(np.std(times)) * 1000,
+    }
 
 
 # ---------------------------------------------------------------------------
@@ -143,6 +189,23 @@ def get_exportable_model(args, rank, world_size):
             .eval()
             .to(DEVICE)
         )
+        # Keep SDPA as a single op so the TRT custom converter (and optional
+        # static KV cache lowering pass) can pattern-match it.
+        register_sdpa.enable_sdpa_converter(args.model, model.config)
+
+    # Column/row slicing divides the head dims by world_size, so both head
+    # counts must be divisible. Small-KV models (e.g. Qwen2.5-0.5B has
+    # num_key_value_heads=2) therefore only support up to TP=2. Fail early with
+    # a clear message instead of a downstream reshape/shape mismatch.
+    assert model.config.num_attention_heads % world_size == 0, (
+        f"num_attention_heads ({model.config.num_attention_heads}) must be "
+        f"divisible by world_size ({world_size})."
+    )
+    assert model.config.num_key_value_heads % world_size == 0, (
+        f"num_key_value_heads ({model.config.num_key_value_heads}) must be "
+        f"divisible by world_size ({world_size}). "
+        f"{args.model} caps out at TP={model.config.num_key_value_heads}."
+    )
 
     # Get the default process group name for NCCL all-reduce.
     default_pg = dist.distributed_c10d._get_default_group()
@@ -179,12 +242,26 @@ def get_exportable_model(args, rank, world_size):
                 _RowParallelLinear(proj, group_name),
             )
 
-    # Patch head counts for the sharded attention
+    # Patch head counts for the sharded attention. Each rank now holds only
+    # num_heads // world_size of the projection outputs, so the per-head reshape
+    # must use the per-rank counts.
+    #
+    # Llama-style attention reads self.num_heads / self.num_key_value_heads for
+    # this reshape, so we set them. Qwen3-style attention instead infers the
+    # head count via a `-1` reshape against self.head_dim (so num_heads may be
+    # absent), and derives the GQA repeat from num_key_value_groups -- a ratio
+    # that's invariant under even sharding. Guard every attribute with hasattr
+    # so the same code is correct for both families.
+    heads_per_rank = model.config.num_attention_heads // world_size
+    kv_heads_per_rank = model.config.num_key_value_heads // world_size
     for layer in model.model.layers:
-        layer.self_attn.num_heads = model.config.num_attention_heads // world_size
-        layer.self_attn.num_key_value_heads = (
-            model.config.num_key_value_heads // world_size
-        )
+        attn = layer.self_attn
+        if hasattr(attn, "num_heads"):
+            attn.num_heads = heads_per_rank
+        if hasattr(attn, "num_key_value_heads"):
+            attn.num_key_value_heads = kv_heads_per_rank
+        if hasattr(attn, "num_key_value_groups"):
+            attn.num_key_value_groups = heads_per_rank // kv_heads_per_rank
 
     logger.info(f"Weights sliced + all-reduce inserted for rank {rank}/{world_size}.")
     return model
@@ -234,12 +311,20 @@ def export_and_save(input_ids, args):
             )
     logger.info("Export succeeded.")
 
+    # Importing these modules registers the static KV cache lowering passes
+    # via @_aten_lowering_pass. Must happen before torch_tensorrt.dynamo.compile.
+    if args.cache == "static_v1":
+        import static_cache_v1  # noqa: F401
+    elif args.cache == "static_v2":
+        import static_cache_v2  # noqa: F401
+
     logger.info("Compiling exported program with TRT (AOT) ...")
-    with (
-        torch_tensorrt.logging.debug()
-        if args.debug
-        else torch.autocast("cuda", dtype=torch.float16)
-    ):
+    # Always run compile under FP16 autocast (matches the verify path below);
+    # additionally enable verbose Torch-TRT logging when --debug is set.
+    with ExitStack() as _compile_stack:
+        _compile_stack.enter_context(torch.autocast("cuda", dtype=torch.float16))
+        if args.debug:
+            _compile_stack.enter_context(torch_tensorrt.logging.debug())
         trt_model = torch_tensorrt.dynamo.compile(
             ep,
             inputs=[
@@ -280,11 +365,19 @@ def export_and_save(input_ids, args):
     logger.info("NCCL communicator eagerly initialized for export verification")
 
     # Verify
-    logger.info("Verifying compiled model ...")
-    with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
-        ref = _extract_logits(model(input_ids, position_ids=position_ids))
-        trt = _extract_logits(trt_model(input_ids, position_ids=position_ids))
-    logger.info(f"Max logit diff: {(ref.float() - trt.float()).abs().max().item():.6f}")
+    if args.cache:
+        # With KV cache, the engine takes (input_ids, position_ids, *kv_cache,
+        # start_idx, end_idx) and the reference model doesn't — logit
+        # comparison requires a separate KV-aware driver. Skip.
+        logger.info("Skipping logit-diff verify (KV cache enabled).")
+    else:
+        logger.info("Verifying compiled model ...")
+        with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16):
+            ref = _extract_logits(model(input_ids, position_ids=position_ids))
+            trt = _extract_logits(trt_model(input_ids, position_ids=position_ids))
+        logger.info(
+            f"Max logit diff: {(ref.float() - trt.float()).abs().max().item():.6f}"
+        )
 
     # Save outside autocast — serialization doesn't need it and retrace=True
     # would fail (execute_engine has no AutocastCUDA kernel for torch.export).
@@ -326,7 +419,8 @@ def load_and_run(input_ids, tokenizer, args):
     logger.info("Engine loaded.")
 
     max_len = input_ids.shape[1] + args.num_tokens
-    loaded_tokens = generate_greedy(
+    gen_fn = _pick_generate_fn(args)
+    loaded_tokens = gen_fn(
         trt_model,
         input_ids.clone(),
         max_len,
@@ -347,48 +441,124 @@ def load_and_run(input_ids, tokenizer, args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
-        description="Llama TP: torch.export → save → load with TRT engines"
+        description="LLM (Llama/Qwen) TP: torch.export → save → load with TRT engines"
+    )
+    parser.add_argument(
+        "--model",
+        default="Qwen/Qwen3-1.7B",
+        help="HF model id. All ungated unless noted. Tested options:\n"
+        "  Qwen/Qwen2.5-0.5B-Instruct  - smallest/fastest smoke test; kv_heads=2 => TP=2 ONLY\n"
+        "  Qwen/Qwen3-1.7B             - default; kv_heads=8 => TP 2/4/8\n"
+        "  meta-llama/Llama-3.2-1B-Instruct - original example (GATED, needs HF token)",
     )
-    parser.add_argument("--model", default="meta-llama/Llama-3.2-1B-Instruct")
     parser.add_argument("--prompt", default="What is tensor parallelism?")
-    parser.add_argument("--num_tokens", type=int, default=64)
-    parser.add_argument("--max_seq_len", type=int, default=128)
+    parser.add_argument("--num_tokens", type=int, default=128)
     parser.add_argument(
         "--mode",
         required=True,
         choices=["export", "load"],
         help="export: AOT compile + save engines | load: load engines + infer",
     )
-    parser.add_argument("--save_dir", default="/tmp/llama_tp_engines")
+    parser.add_argument("--save_dir", default="/tmp/llm_tp_engines")
+    parser.add_argument(
+        "--cache",
+        choices=["", "static_v1", "static_v2"],
+        default="",
+        help="KV cache lowering pass. '' disables cache (full-seq recompute).",
+    )
     parser.add_argument("--debug", action="store_true")
+    parser.add_argument(
+        "--benchmark", action="store_true", help="Measure generation latency"
+    )
+    parser.add_argument(
+        "--iterations", type=int, default=5, help="Benchmark iterations"
+    )
+    parser.add_argument(
+        "--batch_size", type=int, default=1, help="Batch size for benchmarking"
+    )
+    parser.add_argument(
+        "--isl", type=int, default=2048, help="Input sequence length for benchmarking"
+    )
     args = parser.parse_args()
 
     tokenizer = AutoTokenizer.from_pretrained(args.model)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    input_ids = tokenizer(args.prompt, return_tensors="pt")["input_ids"].to(DEVICE)
+    if args.benchmark:
+        input_ids = torch.randint(
+            1, 10000, (args.batch_size, args.isl), dtype=torch.int64
+        ).to(DEVICE)
+    else:
+        input_ids = tokenizer(args.prompt, return_tensors="pt")["input_ids"].to(DEVICE)
     max_len = input_ids.shape[1] + args.num_tokens
+    args.max_seq_len = max_len
 
     trt_model = None
     with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
+        gen_fn = _pick_generate_fn(args)
+
         if args.mode == "export":
             trt_model = export_and_save(input_ids.clone(), args)
 
             logger.info("Running freshly compiled model ...")
-            trt_tokens = generate_greedy(
+            trt_tokens = gen_fn(
                 trt_model,
                 input_ids.clone(),
                 max_len,
                 tokenizer.eos_token_id,
             )
-            if rank == 0:
-                print("\n===== TensorRT-TP (freshly compiled) =====")
-                print(tokenizer.decode(trt_tokens[0], skip_special_tokens=True))
-                sys.stdout.flush()
+            if not args.benchmark:
+                if rank == 0:
+                    print("\n===== TensorRT-TP (freshly compiled) =====")
+                    print(tokenizer.decode(trt_tokens[0], skip_special_tokens=True))
+                    sys.stdout.flush()
+            else:
+                # All ranks must participate in the benchmark loop.
+                use_cache = args.cache in ("static_v1", "static_v2")
+                trt_results = time_generate_split(
+                    trt_model,
+                    input_ids.clone(),
+                    max_len,
+                    tokenizer.eos_token_id,
+                    iterations=args.iterations,
+                    use_cache=use_cache,
+                )
+                if rank == 0:
+                    stats = record_stats_split(
+                        "TensorRT-TP (export)",
+                        trt_results,
+                        "FP16",
+                        batch_size=args.batch_size,
+                    )
+                    print("\n=========TensorRT-TP (export) PERFORMANCE============")
+                    print(stats)
+                    sys.stdout.flush()
 
         elif args.mode == "load":
             trt_model, _ = load_and_run(input_ids, tokenizer, args)
+            if args.benchmark:
+                # All ranks must participate — the decode loop contains NCCL
+                # all-reduce ops that require every rank to call in lockstep.
+                use_cache = args.cache in ("static_v1", "static_v2")
+                trt_results = time_generate_split(
+                    trt_model,
+                    input_ids.clone(),
+                    max_len,
+                    tokenizer.eos_token_id,
+                    iterations=args.iterations,
+                    use_cache=use_cache,
+                )
+                if rank == 0:
+                    stats = record_stats_split(
+                        "TensorRT-TP (load)",
+                        trt_results,
+                        "FP16",
+                        batch_size=args.batch_size,
+                    )
+                    print("\n=========TensorRT-TP (load) PERFORMANCE============")
+                    print(stats)
+                    sys.stdout.flush()
 
     # Delete the TRT engine before destroying the process group — the engine
     # holds a reference to the NCCL communicator and will segfault if NCCL is