chore: import upstream snapshot with attribution

external/get_collections.py

@@ -0,0 +1,90 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+""" Script responsible for generation of a JSON file with list of NeMo collections. """
+
+import argparse
+import importlib
+import json
+import os
+
+import nemo
+from nemo.utils import logging
+
+
+def process_collection(id, col):
+    """Helper function processing the collection.
+
+    Args:
+        id: (short) name of the collection.
+        col: a collection (python module).
+    """
+    return {
+        "id": id,
+        "name": col.__name__,
+        "description": col.__description__,
+        "version": col.__version__,
+        "author": col.__author__,
+    }
+
+
+def main():
+    """Main function generating a JSON file with list of NeMo collections."""
+    # Parse filename.
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--filename', help='Name of the output JSON file', type=str, default="collections.json")
+    args = parser.parse_args()
+
+    # Get collections directory.
+    colletions_dir = os.path.dirname(nemo.collections.__file__)
+    logging.info('Analysing collections in `{}`'.format(colletions_dir))
+
+    # Generate list of NeMo collections - from the list of collection subfolders.
+    collections = {}
+    for sub_dir in os.listdir(colletions_dir):
+        # Skip cache.
+        if sub_dir == "__pycache__":
+            continue
+        # Check if it is a directory.
+        if os.path.isdir(os.path.join(colletions_dir, sub_dir)):
+            collections[sub_dir] = "nemo.collections." + sub_dir
+
+    output_list = []
+    # Iterate over all collections.
+    for key, val in collections.items():
+        # Try to get module specification.
+        module_spec = importlib.util.find_spec(val)
+        if module_spec is None:
+            logging.warning("  * Failed to process `{}`".format(val))
+        else:
+            try:
+                # Import the module from the module specification.
+                module = importlib.util.module_from_spec(module_spec)
+                module_spec.loader.exec_module(module)
+                # Add to list.
+                output_list.append(process_collection(key, module))
+                logging.info("  * Processed `{}`".format(val))
+            except AttributeError:
+                logging.warning("  * Failed to process `{}`".format(val))
+
+    # Export to JSON.
+    with open(args.filename, 'w', encoding='utf-8') as outfile:
+        json.dump(output_list, outfile)
+
+    logging.info('Finshed the analysis, results exported to `{}`.'.format(args.filename))
+
+
+if __name__ == '__main__':
+    main()

external/get_modules.py

@@ -0,0 +1,159 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+""" Script responsible for generation of a JSON file containing list of modules of a given collection. """
+
+import argparse
+import importlib
+import inspect
+import json
+import os
+
+import nemo
+from nemo.utils import logging
+
+
+def process_member(name, obj, module_list):
+    """Helper function processing the passed object and, if ok, adding a record to the module list.
+
+    Args:
+        name: name of the member
+        obj: member (class/function etc.)
+        module_list: list of modules that (probably) will be expanded.
+    """
+    # It is not a class - skip it.
+    if not inspect.isclass(obj):
+        return
+
+    # Check inheritance - we know that all our datasets/modules/losses inherit from Serialization,
+    # Btw. Serialization is also required by this script.
+    if not issubclass(obj, nemo.core.Serialization):
+        return
+
+    logging.info("  * Processing `{}`".format(str(obj)))
+
+    module_list.append(
+        {
+            "name": name,
+            "cls": str(obj),
+            # Temporary solution: mockup arguments.
+            "arguments": [
+                "jasper",
+                "activation",
+                "feat_in",
+                "normalization_mode",
+                "residual_mode",
+                "norm_groups",
+                "conv_mask",
+                "frame_splicing",
+                "init_mode",
+            ],
+            # Temporary solution: mockup input types.
+            "input_types": {
+                "audio_signal": "axes: (batch, dimension, time); elements_type: MelSpectrogramType",
+                "length": "axes: (batch,); elements_type: LengthType",
+            },
+            # Temporary solution: mockup output types.
+            "output_types": {
+                "encoder_output": "axes: (batch, dimension, time); elements_type: AcousticEncodedRepresentation"
+            },
+        }
+    )
+
+
+def main():
+    """Main function analysing the indicated NeMo collection and generating a JSON file with module descriptions."""
+    # Parse filename.
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--collection', help='ID of the collection', type=str)
+    parser.add_argument('--filename', help='Name of the output JSON file', type=str, default="modules.json")
+    args = parser.parse_args()
+
+    # Get collections directory.
+    colletions_dir = os.path.dirname(nemo.collections.__file__)
+    logging.info('Analysing collections in `{}`'.format(colletions_dir))
+
+    # Generate list of NeMo collections - from the list of collection subfolders.
+    collections = {}
+    for sub_dir in os.listdir(colletions_dir):
+        # Skip cache.
+        if sub_dir == "__pycache__":
+            continue
+        # Check if it is a directory.
+        if os.path.isdir(os.path.join(colletions_dir, sub_dir)):
+            collections[sub_dir] = "nemo.collections." + sub_dir
+
+    # Check the collection.
+    if args.collection not in collections.keys():
+        logging.error("Coudn't process the incidated `{}` collection".format(args.collection))
+        logging.info(
+            "Please select one of the existing collections using `--collection [{}]`".format("|".join(collections))
+        )
+        exit(-1)
+
+    # Load the collection specification.
+    collection_spec = importlib.util.find_spec(collections[args.collection])
+    if collection_spec is None:
+        logging.error("Failed to load the `{}` collection".format(val))
+
+    # Import the module from the module specification.
+    collection = importlib.util.module_from_spec(collection_spec)
+    collection_spec.loader.exec_module(collection)
+
+    module_list = []
+    # Iterate over the packages in the indicated collection.
+    logging.info("Analysing the `{}` collection".format(args.collection))
+
+    try:  # Datasets in dataset folder
+        logging.info("Analysing the 'data' package")
+        for name, obj in inspect.getmembers(collection.data):
+            process_member(name, obj, module_list)
+    except AttributeError as e:
+        logging.info("  * No datasets found")
+
+    try:  # Datasets in dataset folder
+        logging.info("Analysing the 'datasets' package")
+        for name, obj in inspect.getmembers(collection.datasets):
+            process_member(name, obj, module_list)
+    except AttributeError as e:
+        logging.info("  * No datasets found")
+
+    try:  # Modules
+        logging.info("Analysing the 'modules' package")
+        for name, obj in inspect.getmembers(collection.modules):
+            process_member(name, obj, module_list)
+    except AttributeError as e:
+        logging.info("  * No modules found")
+
+    try:  # Losses
+        logging.info("Analysing the 'losses' package")
+        for name, obj in inspect.getmembers(collection.losses):
+            process_member(name, obj, module_list)
+    except AttributeError as e:
+        logging.info("  * No losses found")
+
+    # Add prefix - only for default name.
+    filename = args.filename if args.filename != "modules.json" else args.collection + "_" + args.filename
+    # Export to JSON.
+    with open(filename, 'w', encoding='utf-8') as outfile:
+        json.dump(module_list, outfile)
+
+    logging.info(
+        'Finished analysis of the `{}` collection, results exported to `{}`.'.format(args.collection, filename)
+    )
+
+
+if __name__ == '__main__':
+    main()

external/patches/deepep_v1_a100_disable_nvshmem.patch

@@ -0,0 +1,13 @@
+diff --git a/setup.py b/setup.py
+index 6a7f6d4..e2a92f1 100644
+--- a/setup.py
++++ b/setup.py
+@@ -16,7 +16,7 @@ def get_nvshmem_host_lib_name(base_dir):
+ 
+ 
+ if __name__ == '__main__':
+-    disable_nvshmem = False
++    disable_nvshmem = True
+     nvshmem_dir = os.getenv('NVSHMEM_DIR', None)
+     nvshmem_host_lib = 'libnvshmem_host.so'
+     if nvshmem_dir is None:

external/patches/flash_attn_4_sm120_disable_tma_o.patch

@@ -0,0 +1,13 @@
+diff --git a/flash_fwd.py b/flash_fwd.py
+index 7c8cb53..d516e2b 100644
+--- a/flash_fwd.py
++++ b/flash_fwd.py
+@@ -655,7 +655,7 @@ class FlashAttentionForwardSm80:
+         self.num_Q_load_threads = self.num_threads
+         self.num_epilogue_threads = self.num_threads
+         # self.use_tma_O = self.arch >= 90 and mCuSeqlensQ is None
+-        self.use_tma_O = self.arch >= Arch.sm_90
++        self.use_tma_O = self.arch >= Arch.sm_90 and self.arch < Arch.sm_120
+         self._setup_attributes()
+         SharedStorage = self._get_shared_storage_cls()
+         mQ, mK, mV, mO = [assume_tensor_aligned(t) for t in (mQ, mK, mV, mO)]

external/patches/triton-lang_triton_6570_lazy_init.patch

@@ -0,0 +1,66 @@
+/*
+ * Code imported via patch from https://github.com/triton-lang/triton/pull/6570, commit 2afae45951b74785b144151b31e91e6c82b0b02f.
+ * Copyright (c) 2018-2022 Philippe Tillet, OpenAI.
+ * Licensed under the MIT License.
+ */
+
+From 2afae45951b74785b144151b31e91e6c82b0b02f Mon Sep 17 00:00:00 2001
+From: Han Zhu <zhuhan7737@gmail.com>
+Date: Tue, 22 Apr 2025 18:42:23 -0700
+Subject: [PATCH] [autotuner] Lazily initiailize do_bench
+
+---
+diff --git a/a/autotuner.py b/b/autotuner.py
+index 0ee6bea09..b75c5e353 100644
+--- a/a/autotuner.py
++++ b/b/autotuner.py
+@@ -4,6 +4,7 @@ import builtins
+ import os
+ import time
+ import inspect
++from functools import cached_property
+ from typing import Dict, Tuple, List, Optional
+
+ from .jit import KernelInterface
+@@ -94,6 +95,7 @@ class Autotuner(KernelInterface):
+         while not inspect.isfunction(self.base_fn):
+             self.base_fn = self.base_fn.fn
+
++        self._do_bench = do_bench
+         self.num_warmups = warmup
+         self.num_reps = rep
+         self.use_cuda_graph = use_cuda_graph
+@@ -107,7 +109,7 @@ class Autotuner(KernelInterface):
+                           stacklevel=1)
+             if use_cuda_graph:
+                 from ..testing import do_bench_cudagraph
+-                self.do_bench = lambda kernel_call, quantiles: do_bench_cudagraph(
++                self._do_bench = lambda kernel_call, quantiles: do_bench_cudagraph(
+                     kernel_call,
+                     rep=rep if rep is not None else 100,
+                     quantiles=quantiles,
+@@ -115,7 +117,7 @@ class Autotuner(KernelInterface):
+                 return
+
+             import triton.testing
+-            self.do_bench = lambda kernel_call, quantiles: triton.testing.do_bench(
++            self._do_bench = lambda kernel_call, quantiles: triton.testing.do_bench(
+                 kernel_call,
+                 warmup=warmup if warmup is not None else 25,
+                 rep=rep if rep is not None else 100,
+@@ -123,10 +125,11 @@ class Autotuner(KernelInterface):
+             )
+             return
+
+-        if do_bench is None:
+-            self.do_bench = driver.active.get_benchmarker()
+-        else:
+-            self.do_bench = do_bench
++    @cached_property
++    def do_bench(self):
++        if self._do_bench is None:
++            return driver.active.get_benchmarker()
++        return self._do_bench
+
+     def _bench(self, *args, config, **meta):
+         from ..compiler.errors import CompileTimeAssertionFailure

external/patches/trt_llm.patch

@@ -0,0 +1,24 @@
+diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
+index d82bebb73c..b52668a54b 100644
+--- a/cpp/CMakeLists.txt
++++ b/cpp/CMakeLists.txt
+@@ -377,6 +377,7 @@ print(os.path.dirname(torch.__file__),end='');"
+   set(nvtx3_dir ${3RDPARTY_DIR}/NVTX/include)
+   find_package(Torch REQUIRED)
+
++  set(TORCH_CXX_FLAGS "-D_GLIBCXX_USE_CXX11_ABI=1")
+   message(STATUS "TORCH_CXX_FLAGS: ${TORCH_CXX_FLAGS}")
+   add_compile_options(${TORCH_CXX_FLAGS})
+   add_compile_definitions(TORCH_CUDA=1)
+diff --git a/docker/common/install_tensorrt.sh b/docker/common/install_tensorrt.sh
+index 8dda255229..98bdae58f3 100644
+--- a/docker/common/install_tensorrt.sh
++++ b/docker/common/install_tensorrt.sh
+@@ -28,6 +28,7 @@ for i in "$@"; do
+         --CUDNN_VER=?*) CUDNN_VER="${i#*=}";;
+         --NCCL_VER=?*) NCCL_VER="${i#*=}";;
+         --CUBLAS_VER=?*) CUBLAS_VER="${i#*=}";;
++        --NVRTC_VER=?*) NVRTC_VER="${i#*=}";;
+         *) ;;
+     esac
+     shift