chore: import upstream snapshot with attribution

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()