chore: import upstream snapshot with attribution

samples/python/attention_mdtrt/README.md

@@ -0,0 +1,133 @@
+# Multi-Device Attention Inference with TensorRT
+
+This sample demonstrates how to run a self-attention model across multiple GPUs using TensorRT's multi-device inference feature. It covers single-GPU execution as a baseline and multi-GPU execution using MPI for process management and NCCL for GPU-to-GPU communication.
+
+## Introduction
+
+TensorRT supports splitting a single model across multiple GPUs for inference. This is useful when a model is too large to fit on a single GPU, or when you want to reduce latency by parallelizing computation across devices.
+
+TensorRT supports multiple parallelism strategies for multi-device inference, including tensor parallelism (TP) and context parallelism (CP). This sample focuses on **context parallelism (CP)**, where the input sequence is split across GPUs along the sequence dimension. Each GPU processes its portion independently for most operations (linear projections, normalization), but attention requires cross-device communication because every token must attend to every other token. TensorRT handles this communication transparently using collective operations embedded in the model. For a detailed explanation of context parallelism, see [Context Parallelism for Scalable Million-Token Inference](https://arxiv.org/abs/2411.01783).
+
+### How Does it Work?
+
+To run a model on multiple GPUs, the model must be **sharded** (split into pieces that each GPU can execute independently). For context parallelism, sharding means dividing the input sequence across GPUs and inserting communication ops so that each GPU can still compute correct attention over the full sequence.
+
+When a model is sharded for multi-device execution, special **DistCollective** operations are inserted into the ONNX graph:
+
+- **ReduceScatter**: Splits and reduces input data across GPUs. Used at the start to distribute the input sequence.
+- **AllGather**: Collects data from all GPUs into a full tensor. Used before attention so each GPU sees the full K/V tensors, and after the output projection to reconstruct the full output.
+
+These ops are not present in the single-device ONNX model. They are inserted automatically by `polygraphy multi-device shard` using sharding hints that describe how the model should be split.
+
+### What is hint.json?
+
+The `hint.json` file tells the polygraphy sharding tool how to partition the model. The sharding tool reads the single-device ONNX graph, identifies the attention layers and I/O tensors specified in the hints, and inserts DistCollective ops at the appropriate points.
+
+```json
+{
+    "parallelism": "CP",
+    "attention_layers": [
+        {
+            "q": "q_scaled",          // Name of the Q tensor feeding into QK^T matmul
+            "gather_kv": true,        // Gather K/V across devices before attention
+            "gather_q": false         // Q stays local (each GPU has its own chunk)
+        }
+    ],
+    "dist_collectives": {
+        "nb_rank": 2,                 // Number of GPUs to shard across
+        "reduce_op": "max"            // Reduction operation for ReduceScatter
+    },
+    "inputs": [
+        {
+            "name": "input",          // Input tensor name in the ONNX model
+            "seq_len_idx": 0,         // Which dimension is the sequence length
+            "rank": 3                 // Number of dimensions
+        }
+    ],
+    "outputs": [
+        {
+            "name": "output",
+            "seq_len_idx": 0,
+            "rank": 3
+        }
+    ]
+}
+```
+
+For full documentation on the sharding tool and hint format, see the [Polygraphy multi-device documentation](https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy/polygraphy/tools/multi_device).
+
+## Prerequisites
+
+- A machine with multi-GPU
+- `polygraphy` >= 0.49.25 (for `multi-device shard` support)
+
+## Install Dependencies
+
+```bash
+pip3 install -r requirements.txt
+```
+
+## Generating the ONNX Models
+
+### Step 1: Generate the single-device model
+
+```bash
+python3 create_onnx.py --output attention_sd.onnx
+```
+
+This creates a self-attention model with:
+- Input/Output: `(sequence_length, batch_size, 4096)` in float16
+- 32 attention heads, 128-dim per head
+- Q/K/V projections, RMSNorm, scaled dot-product attention, output projection
+
+### Step 2: Shard for multi-device
+
+```bash
+polygraphy multi-device shard attention_sd.onnx -s hint.json -o attention_md.onnx
+```
+
+This inserts DistCollective ops (ReduceScatter, AllGather) into the model based on `hint.json`. The resulting `attention_md.onnx` is designed to run on 2 GPUs.
+
+## Running the Sample
+
+### Single-GPU
+
+```bash
+python3 attention_mdtrt.py \
+  --onnx-path attention_sd.onnx \
+  --sequence-length 56320 \
+  --batch-size 1 \
+  --num-iterations 50
+```
+
+### Multi-GPU (2 GPUs)
+
+```bash
+mpirun -np 2 python3 attention_mdtrt.py \
+  --onnx-path attention_md.onnx \
+  --sequence-length 56320 \
+  --batch-size 1 \
+  --num-iterations 50
+```
+
+### Using a specific libnccl.so
+
+```bash
+LD_PRELOAD=/path/to/libnccl.so mpirun -np 2 python3 attention_mdtrt.py \
+  --onnx-path attention_md.onnx
+```
+
+## License
+
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
+
+## Changelog
+
+April 2026
+Added `create_onnx.py` for ONNX model generation using the GraphSurgeon layer API. Added `--save-output` flag for saving inference output. Updated documentation with DistCollective and sharding explanations.
+
+January 2026
+Initial release of this sample.
+
+## Known Issues
+None

samples/python/attention_mdtrt/attention_mdtrt.py

@@ -0,0 +1,491 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import numpy as np
+import tensorrt as trt
+import torch
+import sys
+import os
+import time
+import argparse
+from cuda.bindings import runtime as cudart
+from ctypes import py_object, pythonapi, c_void_p, c_char_p
+from typing import Optional
+
+try:
+    from mpi4py import MPI
+except ImportError:
+    MPI = None
+
+try:
+    import nccl.core as nccl
+except ImportError:
+    nccl = None
+
+sys.path.insert(1, os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+import common
+
+
+def communicator_to_capsule(comm):
+    """
+    Convert nccl.core.Communicator to PyCapsule for TensorRT compatibility.
+    
+    Args:
+        comm: nccl.core.Communicator instance with .ptr attribute set to ncclComm_t handle
+        
+    Returns:
+        PyCapsule wrapping the communicator pointer, suitable for set_communicator()
+        
+    Raises:
+        ValueError: If comm.ptr is invalid (0 or None), indicating destroyed communicator
+        TypeError: If comm doesn't have a .ptr attribute
+    """
+    # Validate input
+    if comm is None:
+        raise TypeError("Communicator cannot be None")
+    
+    if not hasattr(comm, 'ptr'):
+        raise TypeError(f"Object {type(comm)} does not have 'ptr' attribute. "
+                       "Expected nccl.core.Communicator instance.")
+    
+    # Get the raw pointer from the Communicator object
+    ptr = comm.ptr
+    
+    # Validate that communicator is still alive (ptr != 0)
+    if ptr == 0:
+        raise ValueError("NCCL Communicator has been destroyed (ptr=0). "
+                        "Cannot create capsule for destroyed communicator.")
+    
+    # Convert to PyCapsule using ctypes.pythonapi
+    PyCapsule_New = pythonapi.PyCapsule_New
+    PyCapsule_New.restype = py_object
+    PyCapsule_New.argtypes = [c_void_p, c_char_p, c_void_p]
+    
+    capsule = PyCapsule_New(c_void_p(ptr), b"ncclComm_t", None)
+    
+    if capsule is None:
+        raise RuntimeError("Failed to create PyCapsule from communicator pointer")
+    
+    return capsule
+
+
+def allocate_buffers(engine: trt.ICudaEngine, profile_idx: Optional[int] = None, output_shape: Optional[tuple] = None):
+    """Allocate host and device buffers for TensorRT engine."""
+    inputs = []
+    outputs = []
+    bindings = []
+    tensor_names = [engine.get_tensor_name(i) for i in range(engine.num_io_tensors)]
+    for binding in tensor_names:
+        # Pick out the max shape to allocate enough memory for the binding.
+        shape = engine.get_tensor_shape(binding) if profile_idx is None else engine.get_tensor_profile_shape(binding, profile_idx)[-1]
+        shape_valid = np.all([s >= 0 for s in shape])
+        if not shape_valid and profile_idx is None:
+            raise ValueError(f"Binding {binding} has dynamic shape, " +\
+                "but no profile was specified.")
+        # For dynamic shapes, use fixed output shape
+        if output_shape is not None:
+            shape = output_shape
+        size = trt.volume(shape)
+        trt_type = engine.get_tensor_dtype(binding)
+
+        # Allocate host and device buffers
+        if trt_type == trt.DataType.BF16:
+            dtype = np.dtype(np.uint16)
+            bindingMemory = common.HostDeviceMem(size, dtype)
+        elif trt_type == trt.DataType.HALF:
+            dtype = np.dtype(np.uint16)
+            bindingMemory = common.HostDeviceMem(size, dtype)
+        elif trt_type == trt.DataType.FLOAT:
+            dtype = np.dtype(np.float32)
+            bindingMemory = common.HostDeviceMem(size, dtype)
+        else:
+            try:
+                dtype = np.dtype(trt.nptype(trt_type))
+                bindingMemory = common.HostDeviceMem(size, dtype)
+            except TypeError:
+                size = int(size * trt_type.itemsize)
+                bindingMemory = common.HostDeviceMem(size)
+
+        # Append the device buffer to device bindings.
+        bindings.append(int(bindingMemory.device_ptr))
+
+        # Append to the appropriate list.
+        if engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT:
+            inputs.append(bindingMemory)
+        else:
+            outputs.append(bindingMemory)
+    return inputs, outputs, bindings
+
+
+class AttentionSD:
+    """Base class for Attention model using TensorRT (Single Device)"""
+    
+    def __init__(self, mpi_comm, rank, onnx_path):
+        """
+        Initialize the Attention class
+        
+        Args:
+            mpi_comm: MPI communicator
+            rank: Current instance ID
+            onnx_path: Path to the ONNX model
+        """
+        self.onnx_path = onnx_path
+        self.logger = trt.Logger(trt.Logger.WARNING)
+        self.engine = None
+        self.context = None
+        self.inputs = None
+        self.outputs = None
+        self.bindings = None
+        self.mpi_comm = mpi_comm
+        self.rank = rank
+
+    def setup(self, actual_input_shape, output_shape):
+        """
+        Set up everything before doing inference.
+        """
+        engine_string = self.build_serialized_network()
+        self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(engine_string)
+        if self.engine is None:
+            print("Failed deserializing engine!")
+            exit(-1)
+        print("Succeeded deserializing engine!")
+
+        self.context = self.engine.create_execution_context()
+        
+        # For dynamic shapes, we need to specify the actual input shape we want to use
+        input_name = self.engine.get_tensor_name(0)
+        self.context.set_input_shape(input_name, actual_input_shape)
+        
+        # Allocate buffers
+        self.inputs, self.outputs, self.bindings = allocate_buffers(
+            self.engine, profile_idx=0, output_shape=output_shape
+        )
+        
+        num_io = self.engine.num_io_tensors
+        tensor_names = [self.engine.get_tensor_name(i) for i in range(num_io)]
+        for i in range(num_io):
+            self.context.set_tensor_address(tensor_names[i], self.bindings[i])
+
+    def build_serialized_network(self):
+        """Create and serialize a network from the ONNX model."""
+        # Create builder and empty network
+        builder = trt.Builder(self.logger)
+        network = builder.create_network(flags=1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
+        
+        # Setup parser and parse the ONNX model
+        print(f"Parsing ONNX model from {self.onnx_path}")
+        parser = trt.OnnxParser(network, self.logger)
+
+        with open(self.onnx_path, "rb") as f:
+            if not parser.parse(f.read()):
+                print("Failed to parse ONNX model")
+                for error in range(parser.num_errors):
+                    print(parser.get_error(error))
+                return None
+        
+        # Get input dimensions and data type
+        input_tensor = network.get_input(0)
+        input_shape = input_tensor.shape
+        input_name = input_tensor.name
+        input_dtype = input_tensor.dtype
+        
+        print(f"[Rank {self.rank}] Input shape: {input_shape}")
+        print(f"[Rank {self.rank}] Input name: {input_name}")
+        print(f"[Rank {self.rank}] Input data type: {input_dtype}")
+
+        # Create a builder config
+        config = builder.create_builder_config()
+        
+        config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
+        
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 16 * 1024 * 1024 * 1024)  # 16GB workspace
+        config.set_memory_pool_limit(trt.MemoryPoolType.TACTIC_SHARED_MEMORY, 1 * 1024 * 1024 * 1024)  # 1GB shared mem
+        
+        profile = builder.create_optimization_profile()
+        
+        # Set the shape range for the input tensor
+        min_shape = (1, 1, 4096) 
+        opt_shape = (56320, 1, 4096)
+        max_shape = (56320, 1, 4096)
+        
+        profile.set_shape(input_name, min_shape, opt_shape, max_shape)
+        config.add_optimization_profile(profile)
+        
+        # Build the serialized network
+        serialized_engine = builder.build_serialized_network(network, config)
+        if serialized_engine is None:
+            print(f"[Rank {self.rank}] Failed building serialized engine!")
+            exit(-1)
+        print(f"[Rank {self.rank}] Succeeded building serialized engine!")
+        
+        return serialized_engine
+    
+    def infer(self, input_data, output_shape, num_iterations):
+        """
+        Execute inference on the input data.
+        
+        Args:
+            input_data: Input data for inference
+            output_shape: Expected output shape for reshaping
+            num_iterations: Number of inference iterations for averaging timing results
+            
+        Returns:
+            output_data: List of output tensors
+        """
+        print(f"[Rank {self.rank}] Input shape: {input_data.shape}")
+        
+        # Copy input data to device
+        for input_buffer in self.inputs:
+            common.memcpy_host_to_device(input_buffer.device_ptr, input_data)
+        
+        # Warmup
+        with common.CudaStreamContext() as stream:
+            self.context.execute_async_v3(stream.stream)
+            common.cuda_call(cudart.cudaStreamSynchronize(stream.stream))
+        
+            # Run inference
+            start = time.time()
+            for _ in range(num_iterations):
+                self.context.execute_async_v3(stream.stream)
+                common.cuda_call(cudart.cudaStreamSynchronize(stream.stream))
+            end = time.time()
+        
+            print(f"[Rank {self.rank}] Time spent in TRT attention: {(end-start)/num_iterations * 1000} ms")
+        
+        # Get output
+        output_data = []
+        for output in self.outputs:
+            common.memcpy_device_to_host(output.host, output.device_ptr)
+
+            # Process based on data type
+            if self.engine.get_tensor_dtype(self.engine.get_tensor_name(1)) == trt.DataType.BF16:
+                numpy_output = np.frombuffer(output.host, dtype=np.uint16).reshape(output_shape)
+                torch_output = torch.from_numpy(numpy_output).view(torch.bfloat16)
+                torch_output = torch_output.reshape(output_shape)
+            elif self.engine.get_tensor_dtype(self.engine.get_tensor_name(1)) == trt.DataType.HALF:
+                numpy_output = np.frombuffer(output.host, dtype=np.float16).reshape(output_shape)
+                torch_output = torch.from_numpy(numpy_output)
+            else:
+                numpy_output = np.frombuffer(output.host, dtype=np.float32).reshape(output_shape)
+                torch_output = torch.from_numpy(numpy_output)
+            
+            output_data.append(torch_output)
+                
+        return output_data
+    
+    def cleanup(self):
+        """
+        Free the buffer resources.
+        """
+        common.free_buffers(self.inputs, self.outputs)
+
+
+class AttentionMD(AttentionSD):
+    """Multi-device Attention model using TensorRT with NCCL for communication"""
+    
+    def __init__(self, mpi_comm, num_ranks, rank, onnx_path):
+        """
+        Initialize the multi-device Attention class
+        
+        Args:
+            mpi_comm: MPI communicator
+            num_ranks: Number of instances/devices
+            rank: Current instance ID
+            onnx_path: Path to the ONNX model
+        """
+        super(AttentionMD, self).__init__(mpi_comm, rank, onnx_path)
+        self.num_ranks = num_ranks
+        self.nccl_comm = None
+
+    def setup_multidevice(self, root):
+        """
+        Set up CUDA devices and initialize NCCL communicator.
+        
+        Args:
+            root: Root rank for communication
+        """
+        assert nccl is not None
+        assert root <= self.num_ranks - 1
+        assert self.rank <= self.num_ranks - 1
+
+        num_devices = common.cuda_call(cudart.cudaGetDeviceCount())
+        assert num_devices >= self.num_ranks
+
+        common.cuda_call(cudart.cudaSetDevice(self.rank))
+
+        if self.rank == root:
+            nccl_comm_id = nccl.get_unique_id()
+        else:
+            nccl_comm_id = None
+
+        nccl_comm_id = self.mpi_comm.bcast(nccl_comm_id, root=root)
+        self.nccl_comm = nccl.Communicator.init(nranks=self.num_ranks, rank=self.rank, unique_id=nccl_comm_id)
+
+    def setup(self, actual_input_shape, output_shape, root=0):
+        """
+        Set up the multi-device environment and build/load the engine
+        
+        Args:
+            root: Root rank for communication
+        """
+        self.setup_multidevice(root)
+        
+        # Load or build TRT engine
+        if self.rank == root:
+            engine_bin = bytes(self.build_serialized_network())
+        else:
+            engine_bin = None
+        
+        # Broadcast the serialized engine from root to all ranks
+        engine_bin = self.mpi_comm.bcast(engine_bin, root=root)
+        
+        # Deserialize the engine
+        self.engine = trt.Runtime(self.logger).deserialize_cuda_engine(engine_bin)
+        if self.engine is None:
+            print(f"[Rank {self.rank}] Failed deserializing engine!")
+            exit(-1)
+        print(f"[Rank {self.rank}] Succeeded deserializing engine!")
+        
+        # Create an execution context
+        self.context = self.engine.create_execution_context()
+        
+        # Set the NCCL communicator for multi-device communication
+        capsule = communicator_to_capsule(self.nccl_comm)
+        if not self.context.set_communicator(capsule):
+            print(f"[Rank {self.rank}] Failed to set communicator")
+            exit(-1)
+
+        # For dynamic shapes, we need to specify the actual input shape we want to use
+        input_name = self.engine.get_tensor_name(0)
+        self.context.set_input_shape(input_name, actual_input_shape)
+        
+        # Allocate buffers for local portion of data
+        self.inputs, self.outputs, self.bindings = allocate_buffers(
+            self.engine, profile_idx=0, output_shape=output_shape
+        )
+        
+        num_io = self.engine.num_io_tensors
+        tensor_names = [self.engine.get_tensor_name(i) for i in range(num_io)]
+        for i in range(num_io):
+            self.context.set_tensor_address(tensor_names[i], self.bindings[i])
+
+
+def generate_random_input(sequence_length, batch_size):
+    """Generate random float16 input data with the given shape."""
+    torch.manual_seed(42)
+    torch_input = torch.rand((sequence_length, batch_size, 4096)).to(torch.float16)
+    input_data = np.ascontiguousarray(torch_input.cpu().numpy())
+    return input_data, (sequence_length, batch_size, 4096)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Sample script for Attention MDTRT")
+    parser.add_argument("--onnx-path", type=str, required=True, help="Path to ONNX model")
+    parser.add_argument("--sequence-length", type=int, default=56320, help="Sequence length for input")
+    parser.add_argument("--batch-size", type=int, default=1, help="Batch size")
+    parser.add_argument("--num-iterations", type=int, default=50, help="Number of inference iterations for timing")
+    parser.add_argument("--save-output", type=str, default=None, help="Save output tensor to .npy file (root rank only)")
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    
+    # Initialize MPI if available
+    if MPI is not None:
+        mpi_comm = MPI.COMM_WORLD
+        num_ranks = mpi_comm.Get_size()
+        rank = mpi_comm.Get_rank()
+        root = 0
+    else:
+        # Fallback for single-process execution
+        mpi_comm = None
+        num_ranks = 1
+        rank = 0
+        root = 0
+    
+    actual_input_shape = (args.sequence_length, args.batch_size, 4096)
+    output_shape = (args.sequence_length, args.batch_size, 4096)
+    
+    # Print configuration
+    if rank == root:
+        print(f"[setup] Configuration:")
+        print(f"[setup]   Number of GPUs: {num_ranks}")
+        print(f"[setup]   Sequence Length: {args.sequence_length}")
+        print(f"[setup]   Batch Size: {args.batch_size}")
+        print(f"[setup]   Data Type: float16")
+        print(f"[setup]   Input Shape: {actual_input_shape}")
+        print(f"[setup]   Output Shape: {output_shape}")
+    
+    # Generate random input data with FULL sequence length (only on root rank)
+    if rank == root:
+        input_data, input_shape = generate_random_input(args.sequence_length, args.batch_size)
+        print(f"[Rank {rank}] Generated random input data with shape: {input_shape}")
+        
+        if num_ranks == 1:
+            print(f"[Rank {rank}] Running single-device inference...")
+            try:
+                attention_sd = AttentionSD(mpi_comm, rank, args.onnx_path)
+                attention_sd.setup(actual_input_shape, output_shape)
+                sd_output = attention_sd.infer(input_data, output_shape, args.num_iterations)[0]
+                print(f"[Rank {rank}] Single-device inference completed")
+                print(f"[Rank {rank}] Output shape: {sd_output.shape}")
+                if args.save_output:
+                    np.save(args.save_output, sd_output.float().cpu().numpy())
+                    print(f"[Rank {rank}] Output saved to {args.save_output}")
+                attention_sd.cleanup()
+            except Exception as e:
+                print(f"[Rank {rank}] Error in single-device inference: {e}")
+                sys.exit(1)
+    else:
+        input_data = None
+    
+    # Broadcast full input data to all ranks for multi-device inference
+    if MPI is not None and num_ranks > 1:
+        input_data = mpi_comm.bcast(input_data, root=root)
+    
+    # Run multi-device inference if num_gpus > 1
+    if num_ranks > 1:
+        if MPI is None:
+            print(f"Error: MPI is required for multi-GPU tests but not available. Ensure you run with mpirun.")
+            sys.exit(1)
+        
+        if nccl is None:
+            print(f"Error: nccl is required for multi-GPU tests but not available.")
+            sys.exit(1)
+        
+        print(f"[Rank {rank}] Running multi-device inference...")
+        try:
+            attention_md = AttentionMD(mpi_comm, num_ranks, rank, args.onnx_path)
+            attention_md.setup(actual_input_shape, output_shape, root)
+            md_output = attention_md.infer(input_data, output_shape, args.num_iterations)[0]
+            print(f"[Rank {rank}] Multi-device inference completed")
+            print(f"[Rank {rank}] Output shape: {md_output.shape}")
+            if rank == root and args.save_output:
+                np.save(args.save_output, md_output.float().cpu().numpy())
+                print(f"[Rank {rank}] Output saved to {args.save_output}")
+            attention_md.cleanup()
+        except Exception as e:
+            print(f"[Rank {rank}] Error in multi-device inference: {e}")
+            sys.exit(1)
+    
+    print(f"[Rank {rank}] Test completed successfully!")
+
+
+if __name__ == "__main__":
+    main()

samples/python/attention_mdtrt/create_onnx.py

@@ -0,0 +1,258 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""
+Construct a self-attention ONNX model using the ONNX GraphSurgeon layer API.
+
+The model implements:
+  1. Q/K/V linear projections (MatMul with 4096x4096 weights)
+  2. Reshape to multi-head layout (seq, batch, 4096) -> (seq, batch, 32, 128)
+  3. RMSNorm on Q and K
+  4. Scaled dot-product attention (QK^T / sqrt(d), softmax, attn * V)
+  5. Reshape back and output projection
+
+Input/Output: (sequence_length, batch_size, 4096), float16
+"""
+
+import argparse
+import math
+
+import numpy as np
+import onnx
+import onnx_graphsurgeon as gs
+
+NUM_HEADS = 32
+HEAD_DIM = 128
+HIDDEN_DIM = NUM_HEADS * HEAD_DIM  # 4096
+OPSET = 17
+
+
+# Register ONNX ops as methods on gs.Graph using the layer API.
+# Each returns the output tensor(s) directly for easy chaining.
+
+@gs.Graph.register()
+def matmul(self, a, b):
+    return self.layer(op="MatMul", inputs=[a, b], outputs=["matmul_out"])[0]
+
+
+@gs.Graph.register()
+def transpose(self, a, perm):
+    return self.layer(op="Transpose", inputs=[a], attrs={"perm": perm}, outputs=["transpose_out"])[0]
+
+
+@gs.Graph.register()
+def reshape(self, data, shape):
+    return self.layer(op="Reshape", inputs=[data, shape], attrs={"allowzero": 0}, outputs=["reshape_out"])[0]
+
+
+@gs.Graph.register()
+def softmax(self, a, axis=-1):
+    return self.layer(op="Softmax", inputs=[a], attrs={"axis": axis}, outputs=["softmax_out"])[0]
+
+
+@gs.Graph.register()
+def cast(self, a, to):
+    return self.layer(op="Cast", inputs=[a], attrs={"to": to}, outputs=["cast_out"])[0]
+
+
+@gs.Graph.register()
+def sqrt(self, a):
+    return self.layer(op="Sqrt", inputs=[a], outputs=["sqrt_out"])[0]
+
+
+@gs.Graph.register()
+def add(self, a, b):
+    return self.layer(op="Add", inputs=[a, b], outputs=["add_out"])[0]
+
+
+@gs.Graph.register()
+def mul(self, a, b):
+    return self.layer(op="Mul", inputs=[a, b], outputs=["mul_out"])[0]
+
+
+@gs.Graph.register()
+def div(self, a, b):
+    return self.layer(op="Div", inputs=[a, b], outputs=["div_out"])[0]
+
+
+@gs.Graph.register()
+def pow(self, a, b):
+    return self.layer(op="Pow", inputs=[a, b], outputs=["pow_out"])[0]
+
+
+@gs.Graph.register()
+def reduce_mean(self, a, axes, keepdims=1):
+    return self.layer(
+        op="ReduceMean", inputs=[a], attrs={"axes": axes, "keepdims": keepdims},
+        outputs=["reduce_mean_out"],
+    )[0]
+
+
+@gs.Graph.register()
+def shape_op(self, a):
+    return self.layer(op="Shape", inputs=[a], outputs=["shape_out"])[0]
+
+
+@gs.Graph.register()
+def gather(self, data, indices):
+    return self.layer(op="Gather", inputs=[data, indices], attrs={"axis": 0}, outputs=["gather_out"])[0]
+
+
+@gs.Graph.register()
+def unsqueeze(self, a, axes):
+    return self.layer(op="Unsqueeze", inputs=[a, axes], outputs=["unsqueeze_out"])[0]
+
+
+@gs.Graph.register()
+def concat(self, inputs, axis=0):
+    return self.layer(op="Concat", inputs=inputs, attrs={"axis": axis}, outputs=["concat_out"])[0]
+
+
+def build_attention_graph():
+    """Build the full self-attention ONNX graph."""
+    rng = np.random.default_rng(42)
+    graph = gs.Graph(opset=OPSET)
+
+    def fp16_weights(shape):
+        return rng.standard_normal(shape).astype(np.float16)
+
+    def fp32_scalar(val):
+        return np.array([val], dtype=np.float32)
+
+    axes_0 = np.array([0], dtype=np.int64)
+
+    # Input: (seq, batch, 4096) fp16
+    graph_input = gs.Variable("input", dtype=np.float16, shape=["sequence_length", "batch_size", HIDDEN_DIM])
+    graph.inputs = [graph_input]
+
+    # Q/K/V projections
+    q_proj = graph.matmul(graph_input, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
+    k_proj = graph.matmul(graph_input, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
+    v_proj = graph.matmul(graph_input, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
+
+    # Dynamic reshape: (seq, batch, 4096) -> (seq, batch, 32, 128)
+    # Build target shape [seq_dim, batch_dim, 32, 128] from input shape
+    def reshape_to_heads(proj):
+        inp_shape = graph.shape_op(proj)
+        seq_dim = graph.unsqueeze(graph.gather(inp_shape, np.array(0, dtype=np.int64)), axes_0)
+        batch_dim = graph.unsqueeze(graph.gather(inp_shape, np.array(1, dtype=np.int64)), axes_0)
+        target_shape = graph.concat([
+            seq_dim, batch_dim,
+            np.array([NUM_HEADS], dtype=np.int64),
+            np.array([HEAD_DIM], dtype=np.int64),
+        ])
+        return graph.reshape(proj, target_shape)
+
+    q_4d = reshape_to_heads(q_proj)
+    k_4d = reshape_to_heads(k_proj)
+    v_4d = reshape_to_heads(v_proj)
+
+    # RMSNorm: x * rsqrt(mean(x^2) + eps) * weight
+    def rmsnorm(x):
+        x_fp32 = graph.cast(x, onnx.TensorProto.FLOAT)
+        sq = graph.pow(x_fp32, fp32_scalar(2.0))
+        mean = graph.reduce_mean(sq, axes=[-1])
+        rms = graph.sqrt(graph.add(mean, fp32_scalar(1e-6)))
+        inv_rms = graph.div(fp32_scalar(1.0), rms)
+        normed = graph.mul(x_fp32, inv_rms)
+        normed_fp16 = graph.cast(normed, onnx.TensorProto.FLOAT16)
+        weight = rng.standard_normal((1, 1, 1, HEAD_DIM)).astype(np.float16)
+        return graph.mul(weight, normed_fp16)
+
+    q_norm = rmsnorm(q_4d)
+    k_norm = rmsnorm(k_4d)
+
+    # Transpose to attention layout: (seq, batch, heads, hdim) -> (batch, heads, seq, hdim)
+    q_attn = graph.transpose(q_norm, perm=[1, 2, 0, 3])
+    k_attn = graph.transpose(k_norm, perm=[1, 2, 0, 3])
+    v_attn = graph.transpose(v_4d, perm=[1, 2, 0, 3])
+
+    # Dynamic reshape Q/K/V to (batch, heads, -1, hdim) using shape extraction
+    def reshape_attn(x):
+        s = graph.shape_op(x)
+        batch = graph.unsqueeze(graph.gather(s, np.array(0, dtype=np.int64)), axes_0)
+        heads = graph.unsqueeze(graph.gather(s, np.array(1, dtype=np.int64)), axes_0)
+        hdim = graph.unsqueeze(graph.gather(s, np.array(3, dtype=np.int64)), axes_0)
+        target = graph.concat([batch, heads, np.array([-1], dtype=np.int64), hdim])
+        return graph.reshape(x, target)
+
+    q_r = reshape_attn(q_attn)
+    k_r = reshape_attn(k_attn)
+    v_r = reshape_attn(v_attn)
+
+    # Scale: split sqrt(1/sqrt(head_dim)) across Q and K
+    scale_val = math.sqrt(math.sqrt(1.0 / HEAD_DIM))
+    scale_fp16 = np.array([scale_val], dtype=np.float16)
+    q_scaled = graph.mul(q_r, scale_fp16)
+    q_scaled.name = "q_scaled"
+
+    # Transpose K: (batch, heads, seq, hdim) -> (batch, heads, hdim, seq)
+    k_t = graph.transpose(k_r, perm=[0, 1, 3, 2])
+    k_scaled = graph.mul(k_t, scale_fp16)
+
+    # QK^T -> Softmax -> Attn*V
+    qk = graph.matmul(q_scaled, k_scaled)
+    attn_weights = graph.softmax(qk, axis=-1)
+    attn_out = graph.matmul(attn_weights, v_r)
+
+    # Reshape back: (batch, heads, seq, hdim) -> (seq, batch, 4096)
+    attn_t = graph.transpose(attn_out, perm=[2, 0, 1, 3])
+    attn_shape = graph.shape_op(attn_t)
+    seq_dim = graph.unsqueeze(graph.gather(attn_shape, np.array(0, dtype=np.int64)), axes_0)
+    batch_dim = graph.unsqueeze(graph.gather(attn_shape, np.array(1, dtype=np.int64)), axes_0)
+    # Compute heads * hdim
+    heads_dim = graph.gather(attn_shape, np.array(2, dtype=np.int64))
+    hdim_dim = graph.gather(attn_shape, np.array(3, dtype=np.int64))
+    hidden = graph.unsqueeze(graph.mul(heads_dim, hdim_dim), axes_0)
+    flat_shape = graph.concat([seq_dim, batch_dim, hidden])
+    attn_flat = graph.reshape(attn_t, flat_shape)
+
+    # Output projection
+    output = graph.matmul(attn_flat, fp16_weights((HIDDEN_DIM, HIDDEN_DIM)))
+    output.name = "output"
+    output.dtype = np.float16
+    output.shape = ["sequence_length", "batch_size", HIDDEN_DIM]
+    graph.outputs = [output]
+
+    graph.cleanup().toposort()
+    model = gs.export_onnx(graph)
+    model.ir_version = 8
+
+    return model
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate attention_sd.onnx model using the ONNX GraphSurgeon layer API"
+    )
+    parser.add_argument(
+        "--output", type=str, default="attention_sd.onnx",
+        help="Output ONNX file path (default: attention_sd.onnx)",
+    )
+    args = parser.parse_args()
+
+    model = build_attention_graph()
+    onnx.save(model, args.output)
+
+    print(f"Saved model to {args.output}")
+    print(f"  Nodes: {len(model.graph.node)}")
+    print(f"  Initializers: {len(model.graph.initializer)}")
+    print(f"  Opset: {model.opset_import[0].version}")
+
+
+if __name__ == "__main__":
+    main()

samples/python/attention_mdtrt/hint.json

@@ -0,0 +1,40 @@
+{
+    "parallelism": "CP",
+    "attention_layers": [
+        {
+            "q": "q_scaled",
+            "gather_kv": true,
+            "gather_q": false,
+            "replace": null,
+            "polygraphy_class": "AttentionLayerHint"
+        }
+    ],
+    "dist_collectives": {
+        "group_size": 0,
+        "root": -1,
+        "nb_rank": 2,
+        "reduce_op": "max",
+        "groups": [],
+        "polygraphy_class": "DistCollective"
+    },
+    "inputs": [
+        {
+            "name": "input",
+            "seq_len_idx": 0,
+            "rank": 3,
+            "polygraphy_class": "ShardTensor"
+        }
+    ],
+    "outputs": [
+        {
+            "name": "output",
+            "seq_len_idx": 0,
+            "rank": 3,
+            "polygraphy_class": "ShardTensor"
+        }
+    ],
+    "k_seq_len_idx": 3,
+    "v_seq_len_idx": 2,
+    "kv_rank": 4,
+    "polygraphy_class": "ShardHints"
+}

samples/python/attention_mdtrt/requirements.txt

@@ -0,0 +1,7 @@
+numpy==1.26.4
+cuda-bindings
+mpi4py==4.1.1
+nccl4py==0.1.1
+onnx_graphsurgeon
+polygraphy>=0.49.25
+torch