From 7143ebf93b9c865f86bd5faf227de9ea55a937b5 Mon Sep 17 00:00:00 2001
From: Onur Yilmaz <oyilmaz@nvidia.com>
Date: Thu, 11 Jun 2026 15:09:44 -0400
Subject: [PATCH 1/5] Addressing new inference updates in mcore

Signed-off-by: Onur Yilmaz <oyilmaz@nvidia.com>
---
 nemo_deploy/llm/inference/inference_base.py   | 74 ++++++++-----------
 nemo_deploy/llm/megatronllm_deployable.py     | 34 ++++-----
 .../functional_tests/utils/run_nemo_export.py |  6 +-
 .../deploy/test_etp_sequence_parallel.py      | 50 ++++---------
 .../unit_tests/deploy/test_inference_base.py  | 61 ++++++---------
 .../test_megatron_multimodal_deployable.py    | 14 ++--
 .../deploy/test_megatronllm_deployable.py     | 51 +++++++------
 7 files changed, 119 insertions(+), 171 deletions(-)

diff --git a/nemo_deploy/llm/inference/inference_base.py b/nemo_deploy/llm/inference/inference_base.py
index 9c8681441..b2f16d4fd 100644
--- a/nemo_deploy/llm/inference/inference_base.py
+++ b/nemo_deploy/llm/inference/inference_base.py
@@ -27,14 +27,8 @@
     get_default_load_sharded_strategy,
 )
 from megatron.core.dist_checkpointing.validation import StrictHandling
-from megatron.core.inference.contexts.static_context import StaticInferenceContext
-from megatron.core.inference.engines.mcore_engine import MCoreEngine
-from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
-    GPTInferenceWrapper,
-)
-from megatron.core.inference.text_generation_controllers.text_generation_controller import (
-    TextGenerationController,
-)
+from megatron.core.inference.apis import MegatronLLM
+from megatron.core.inference.config import InferenceConfig
 from megatron.core.transformer.enums import AttnBackend
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.transformer.transformer_config import MLATransformerConfig
@@ -390,27 +384,24 @@ def setup_model_and_tokenizer_for_inference(
 
 
 class MCoreEngineWithCleanup:
-    """Wrapper around MCoreEngine that ensures proper cleanup of distributed resources.
+    """Wrapper around MegatronLLM that ensures proper cleanup of distributed resources.
 
-    This class delegates all operations to the underlying MCoreEngine while ensuring that
-    distributed resources are properly cleaned up when the engine is destroyed.
+    This class delegates all operations to the underlying MegatronLLM engine while ensuring
+    that distributed resources are properly cleaned up when the engine is destroyed.
     """
 
     def __init__(
         self,
-        mcore_engine: MCoreEngine,
-        model_inference_wrapper: GPTInferenceWrapper,
+        llm: MegatronLLM,
         tokenizer: Union[MCoreTokenizerWrappper, MegatronTokenizer],
     ):
         """Initialize the MCoreEngineWithCleanup.
 
         Args:
-            mcore_engine (MCoreEngine): The underlying MCoreEngine instance
-            model_inference_wrapper (GPTInferenceWrapper): The model inference wrapper
+            llm (MegatronLLM): The underlying MegatronLLM instance
             tokenizer (Union[MCoreTokenizerWrappper, MegatronTokenizer]): The tokenizer instance
         """
-        self.mcore_engine = mcore_engine
-        self.model_inference_wrapper = model_inference_wrapper
+        self.mcore_engine = llm
         self.tokenizer = tokenizer
 
     def __del__(self):
@@ -446,8 +437,8 @@ def create_mcore_engine(
     buffer_size_gb: float = 10.0,
     legacy_model_format: bool = False,
     **model_config_kwargs,
-) -> Tuple[MCoreEngineWithCleanup, GPTInferenceWrapper, Union[MCoreTokenizerWrappper, MegatronTokenizer]]:
-    """Set up the model, tokenizer and MCoreEngine for inference.
+) -> Tuple[MCoreEngineWithCleanup, Union[MCoreTokenizerWrappper, MegatronTokenizer]]:
+    """Set up the model, tokenizer and MegatronLLM engine for inference.
 
     Args:
         path (Path): Path to the checkpoint file
@@ -455,7 +446,7 @@ def create_mcore_engine(
         inference_batch_times_seqlen_threshold (int): Threshold for batch size times sequence length
         inference_max_seq_length (int): Maximum sequence length for inference
         max_batch_size (int): Maximum batch size for inference
-        random_seed (Optional[int]): Random seed for reproducibility
+        random_seed (Optional[int]): Random seed for reproducibility (set globally during init)
         tensor_model_parallel_size (Optional[int]): Size of tensor model parallelism
         pipeline_model_parallel_size (Optional[int]): Size of pipeline model parallelism
         context_parallel_size (Optional[int]): Size of context parallelism
@@ -466,11 +457,10 @@ def create_mcore_engine(
         model_type (str): Type of model to load (default: "gpt")
         model_format (str): Format of model to load (default: "nemo")
         micro_batch_size (Optional[int]): Micro batch size for model execution
-        legacy_model_format (bool): Whether to use the legacy StaticInferenceEngine path in MCoreEngine (default: False)
+        legacy_model_format (bool): Deprecated; no longer used (DynamicInferenceEngine is always used)
     Returns:
-        Tuple[MCoreEngineWithCleanup, GPTInferenceWrapper, Union[MCoreTokenizerWrappper, MegatronTokenizer]]: Tuple containing:
+        Tuple[MCoreEngineWithCleanup, Union[MCoreTokenizerWrappper, MegatronTokenizer]]: Tuple containing:
             - MCoreEngineWithCleanup: Engine for text generation with proper cleanup
-            - GPTInferenceWrapper: Inference-wrapped model
             - Union[MCoreTokenizerWrappper, MegatronTokenizer]: Tokenizer instance
     """
     # Default to 1 for any parallelism dimension that's None
@@ -512,34 +502,32 @@ def create_mcore_engine(
     else:
         raise ValueError(f"Model format {model_format} not supported.")
 
-    # MLA models require block_size_tokens=64 for the dynamic engine, which is not
-    # configurable in the current Megatron-LM version. Fall back to the legacy static
-    # engine so MLA inference works correctly without touching Megatron-LM.
+    model.eval()
+
+    # MLA models require block_size_tokens=64 for correct KV cache operation with the
+    # dynamic inference engine. Set the attention backend to flash if not already set.
+    block_size_tokens = 256
     model_config = getattr(model, "config", None)
     if isinstance(model_config, MLATransformerConfig):
-        legacy_model_format = True
-        # The legacy static engine requires an explicit attention backend.
-        # MLA models use flash attention (attention_mask is handled internally).
+        block_size_tokens = 64
         if not model_config.attention_backend:
             model_config.attention_backend = AttnBackend.flash
 
-    inference_context = StaticInferenceContext(
-        max_batch_size=max_batch_size,
+    inference_config = InferenceConfig(
         max_sequence_length=inference_max_seq_length,
+        buffer_size_gb=int(buffer_size_gb),
+        max_requests=max_batch_size,
+        block_size_tokens=block_size_tokens,
+        materialize_only_last_token_logits=True,
     )
-    model_inference_wrapper = GPTInferenceWrapper(model, inference_context)
-    text_generation_controller = TextGenerationController(
-        inference_wrapped_model=model_inference_wrapper, tokenizer=tokenizer
-    )
-    mcore_engine = MCoreEngine(
-        text_generation_controller=text_generation_controller,
-        max_batch_size=max_batch_size,
-        random_seed=random_seed,
-        buffer_size_gb=buffer_size_gb,
-        legacy=legacy_model_format,
+
+    llm = MegatronLLM(
+        model=model,
+        tokenizer=tokenizer,
+        inference_config=inference_config,
     )
 
     # Wrap the engine to ensure cleanup
-    wrapped_engine = MCoreEngineWithCleanup(mcore_engine, model_inference_wrapper, tokenizer)
+    wrapped_engine = MCoreEngineWithCleanup(llm, tokenizer)
 
-    return wrapped_engine, model_inference_wrapper, tokenizer
+    return wrapped_engine, tokenizer
diff --git a/nemo_deploy/llm/megatronllm_deployable.py b/nemo_deploy/llm/megatronllm_deployable.py
index ccbeddb12..10910ac90 100755
--- a/nemo_deploy/llm/megatronllm_deployable.py
+++ b/nemo_deploy/llm/megatronllm_deployable.py
@@ -21,8 +21,8 @@
 import torch
 import torch.distributed
 from jinja2 import Template
-from megatron.core.inference.common_inference_params import CommonInferenceParams
-from megatron.core.inference.inference_request import InferenceRequest
+from megatron.core.inference.inference_request import DynamicInferenceRequest
+from megatron.core.inference.sampling_params import SamplingParams
 
 from nemo_deploy import ITritonDeployable
 from nemo_deploy.llm.inference.inference_base import create_mcore_engine
@@ -113,7 +113,7 @@ def __init__(
         if model_type not in ["gpt", "mamba"]:
             raise ValueError(f"Model type {model_type} not supported for Megatron models.")
 
-        self.mcore_engine, self.inference_wrapped_model, self.mcore_tokenizer = create_mcore_engine(
+        self.mcore_engine, self.mcore_tokenizer = create_mcore_engine(
             num_devices=num_devices,
             num_nodes=num_nodes,
             path=Path(megatron_checkpoint_filepath),
@@ -144,18 +144,18 @@ def __init__(
     def generate(
         self,
         prompts: List[str],
-        inference_params: Optional[CommonInferenceParams] = None,
-    ) -> List[InferenceRequest]:
+        inference_params: Optional[SamplingParams] = None,
+    ) -> List[DynamicInferenceRequest]:
         """Generates text based on the provided input prompts.
 
         Args:
             prompts (List[str]): A list of input strings.
-            inference_params (Optional[CommonInferenceParams]): Parameters for controlling the inference process.
+            inference_params (Optional[SamplingParams]): Parameters for controlling the inference process.
 
         Returns:
-            List[InferenceRequest]: A list containing the generated results.
+            List[DynamicInferenceRequest]: A list containing the generated results.
         """
-        inference_params = inference_params or CommonInferenceParams()
+        inference_params = inference_params or SamplingParams()
 
         # Store the original number of prompts
         orig_num_prompts = len(prompts)
@@ -173,8 +173,7 @@ def generate(
 
             results = self.mcore_engine.generate(
                 prompts=padded_prompts,
-                add_BOS=False,
-                common_inference_params=inference_params,
+                sampling_params=inference_params,
             )
 
             # Only return results for the original prompts
@@ -182,8 +181,7 @@ def generate(
         else:
             results = self.mcore_engine.generate(
                 prompts=prompts,
-                add_BOS=False,
-                common_inference_params=inference_params,
+                sampling_params=inference_params,
             )
             return list(results)
 
@@ -198,7 +196,7 @@ def generate_other_ranks(self):
                     data=[None], src=0
                 )
 
-                inference_params = CommonInferenceParams(
+                inference_params = SamplingParams(
                     temperature=temperature,
                     top_k=int(top_k),
                     top_p=float(top_p),
@@ -208,7 +206,7 @@ def generate_other_ranks(self):
                 )
 
                 if log_probs:
-                    dynamic_engine = getattr(self.mcore_engine, "dynamic_engine", None)
+                    dynamic_engine = getattr(self.mcore_engine, "engine", None)
                     if dynamic_engine is not None:
                         dynamic_engine.materialize_only_last_token_logits = False
                         dynamic_engine.context.config.materialize_only_last_token_logits = False
@@ -419,15 +417,15 @@ def _infer_fn(
                 )
 
         # cast top_k,top_p to native int, float since typecheck assert statements added in MCore0.13 error otherwise
-        # return_prompt_top_n_logprobs returns top_logprobs for prompt tokens too when top_logprobs>0.
-        inference_params = CommonInferenceParams(
+        # skip_prompt_log_probs=False (default) includes prompt tokens in top-N logprobs when top_logprobs>0.
+        inference_params = SamplingParams(
             temperature=temperature,
             top_k=int(top_k),
             top_p=float(top_p),
             num_tokens_to_generate=num_tokens_to_generate,
             return_log_probs=log_probs,
             top_n_logprobs=top_logprobs,
-            return_prompt_top_n_logprobs=bool(top_logprobs),
+            skip_prompt_log_probs=not bool(top_logprobs),
             stop_words=stop_words,
         )
 
@@ -436,7 +434,7 @@ def _infer_fn(
         # (prompt log probs are required for logprob eval benchmarks).
         # Toggle it on both the engine and the context config (controls the
         # model forward pass and log prob calculations).
-        dynamic_engine = getattr(self.mcore_engine, "dynamic_engine", None)
+        dynamic_engine = getattr(self.mcore_engine, "engine", None)
         needs_all_logits = log_probs or bool(top_logprobs)
         if dynamic_engine is not None and needs_all_logits:
             dynamic_engine.materialize_only_last_token_logits = False
diff --git a/tests/functional_tests/utils/run_nemo_export.py b/tests/functional_tests/utils/run_nemo_export.py
index ce6938299..521be9255 100644
--- a/tests/functional_tests/utils/run_nemo_export.py
+++ b/tests/functional_tests/utils/run_nemo_export.py
@@ -35,13 +35,13 @@
 
 in_framework_supported = True
 try:
-    from megatron.core.inference.common_inference_params import CommonInferenceParams
+    from megatron.core.inference.sampling_params import SamplingParams
 
     from nemo_deploy.llm import NemoQueryLLMPyTorch
     from nemo_deploy.llm.megatronllm_deployable import MegatronLLMDeployable
 except Exception as e:
     LOGGER.warning(
-        "Cannot import MegatronLLMDeployable class, or NemoQueryLLMPyTorch, or CommonInferenceParams, "
+        "Cannot import MegatronLLMDeployable class, or NemoQueryLLMPyTorch, or SamplingParams, "
         f"in-framework inference will not be available. Reason: {type(e).__name__}: {e}"
     )
     in_framework_supported = False
@@ -98,7 +98,7 @@ def get_accuracy_with_lambada(model, nq, lora_uids, test_data_path, use_vllm: bo
                 if in_framework_supported and isinstance(model, MegatronLLMDeployable):
                     model_output = model.generate(
                         prompts=[prompt],
-                        inference_params=CommonInferenceParams(
+                        inference_params=SamplingParams(
                             temperature=0.1,
                             top_k=1,
                             top_p=0.0,
diff --git a/tests/unit_tests/deploy/test_etp_sequence_parallel.py b/tests/unit_tests/deploy/test_etp_sequence_parallel.py
index 732f2d067..6cdc1069f 100644
--- a/tests/unit_tests/deploy/test_etp_sequence_parallel.py
+++ b/tests/unit_tests/deploy/test_etp_sequence_parallel.py
@@ -389,19 +389,14 @@ class TestCreateMcoreEngineETPSequenceParallel(unittest.TestCase):
     """Tests that create_mcore_engine handles ETP/SP defaults and passes them down."""
 
     @patch("nemo_deploy.llm.inference.inference_base.setup_model_and_tokenizer_for_inference")
-    @patch("nemo_deploy.llm.inference.inference_base.MCoreEngine")
+    @patch("nemo_deploy.llm.inference.inference_base.MegatronLLM")
     @patch("nemo_deploy.llm.inference.inference_base.MCoreEngineWithCleanup")
-    @patch("nemo_deploy.llm.inference.inference_base.GPTInferenceWrapper")
-    @patch("nemo_deploy.llm.inference.inference_base.StaticInferenceContext")
-    @patch("nemo_deploy.llm.inference.inference_base.TextGenerationController")
-    def test_etp_defaults_to_1_when_none(
-        self, mock_tgc, mock_ctx, mock_wrapper, mock_cleanup, mock_engine_cls, mock_setup
-    ):
+    def test_etp_defaults_to_1_when_none(self, mock_cleanup, mock_llm_cls, mock_setup):
         """expert_tensor_parallel_size=None is normalised to 1 before forwarding."""
         from nemo_deploy.llm.inference.inference_base import create_mcore_engine
 
         mock_setup.return_value = ([MagicMock()], MagicMock())
-        mock_engine_cls.return_value = MagicMock()
+        mock_llm_cls.return_value = MagicMock()
         mock_cleanup.return_value = MagicMock()
 
         create_mcore_engine(path=Path("/fake"), model_format="nemo", expert_tensor_parallel_size=None)
@@ -410,19 +405,14 @@ def test_etp_defaults_to_1_when_none(
         assert kwargs["expert_tensor_parallel_size"] == 1
 
     @patch("nemo_deploy.llm.inference.inference_base.setup_model_and_tokenizer_for_inference")
-    @patch("nemo_deploy.llm.inference.inference_base.MCoreEngine")
+    @patch("nemo_deploy.llm.inference.inference_base.MegatronLLM")
     @patch("nemo_deploy.llm.inference.inference_base.MCoreEngineWithCleanup")
-    @patch("nemo_deploy.llm.inference.inference_base.GPTInferenceWrapper")
-    @patch("nemo_deploy.llm.inference.inference_base.StaticInferenceContext")
-    @patch("nemo_deploy.llm.inference.inference_base.TextGenerationController")
-    def test_sp_defaults_to_1_when_none(
-        self, mock_tgc, mock_ctx, mock_wrapper, mock_cleanup, mock_engine_cls, mock_setup
-    ):
+    def test_sp_defaults_to_1_when_none(self, mock_cleanup, mock_llm_cls, mock_setup):
         """sequence_parallel=None is normalised to 1 before forwarding."""
         from nemo_deploy.llm.inference.inference_base import create_mcore_engine
 
         mock_setup.return_value = ([MagicMock()], MagicMock())
-        mock_engine_cls.return_value = MagicMock()
+        mock_llm_cls.return_value = MagicMock()
         mock_cleanup.return_value = MagicMock()
 
         create_mcore_engine(path=Path("/fake"), model_format="nemo", sequence_parallel=None)
@@ -431,19 +421,14 @@ def test_sp_defaults_to_1_when_none(
         assert kwargs["sequence_parallel"] == 1
 
     @patch("nemo_deploy.llm.inference.inference_base.setup_model_and_tokenizer_for_inference")
-    @patch("nemo_deploy.llm.inference.inference_base.MCoreEngine")
+    @patch("nemo_deploy.llm.inference.inference_base.MegatronLLM")
     @patch("nemo_deploy.llm.inference.inference_base.MCoreEngineWithCleanup")
-    @patch("nemo_deploy.llm.inference.inference_base.GPTInferenceWrapper")
-    @patch("nemo_deploy.llm.inference.inference_base.StaticInferenceContext")
-    @patch("nemo_deploy.llm.inference.inference_base.TextGenerationController")
-    def test_explicit_etp_passed_through(
-        self, mock_tgc, mock_ctx, mock_wrapper, mock_cleanup, mock_engine_cls, mock_setup
-    ):
+    def test_explicit_etp_passed_through(self, mock_cleanup, mock_llm_cls, mock_setup):
         """An explicit expert_tensor_parallel_size value is forwarded unchanged."""
         from nemo_deploy.llm.inference.inference_base import create_mcore_engine
 
         mock_setup.return_value = ([MagicMock()], MagicMock())
-        mock_engine_cls.return_value = MagicMock()
+        mock_llm_cls.return_value = MagicMock()
         mock_cleanup.return_value = MagicMock()
 
         create_mcore_engine(path=Path("/fake"), model_format="nemo", expert_tensor_parallel_size=4)
@@ -452,19 +437,14 @@ def test_explicit_etp_passed_through(
         assert kwargs["expert_tensor_parallel_size"] == 4
 
     @patch("nemo_deploy.llm.inference.inference_base.setup_model_and_tokenizer_for_inference")
-    @patch("nemo_deploy.llm.inference.inference_base.MCoreEngine")
+    @patch("nemo_deploy.llm.inference.inference_base.MegatronLLM")
     @patch("nemo_deploy.llm.inference.inference_base.MCoreEngineWithCleanup")
-    @patch("nemo_deploy.llm.inference.inference_base.GPTInferenceWrapper")
-    @patch("nemo_deploy.llm.inference.inference_base.StaticInferenceContext")
-    @patch("nemo_deploy.llm.inference.inference_base.TextGenerationController")
-    def test_explicit_sp_passed_through(
-        self, mock_tgc, mock_ctx, mock_wrapper, mock_cleanup, mock_engine_cls, mock_setup
-    ):
+    def test_explicit_sp_passed_through(self, mock_cleanup, mock_llm_cls, mock_setup):
         """An explicit sequence_parallel=True value is forwarded unchanged."""
         from nemo_deploy.llm.inference.inference_base import create_mcore_engine
 
         mock_setup.return_value = ([MagicMock()], MagicMock())
-        mock_engine_cls.return_value = MagicMock()
+        mock_llm_cls.return_value = MagicMock()
         mock_cleanup.return_value = MagicMock()
 
         create_mcore_engine(path=Path("/fake"), model_format="nemo", sequence_parallel=True)
@@ -487,7 +467,7 @@ def test_expert_tensor_parallel_size_forwarded(self, mock_create):
         """expert_tensor_parallel_size is forwarded to create_mcore_engine."""
         from nemo_deploy.llm.megatronllm_deployable import MegatronLLMDeployable
 
-        mock_create.return_value = (MagicMock(), MagicMock(), MagicMock())
+        mock_create.return_value = (MagicMock(), MagicMock())
 
         MegatronLLMDeployable(
             megatron_checkpoint_filepath="model.ckpt",
@@ -503,7 +483,7 @@ def test_sequence_parallel_forwarded(self, mock_create):
         """sequence_parallel is forwarded to create_mcore_engine."""
         from nemo_deploy.llm.megatronllm_deployable import MegatronLLMDeployable
 
-        mock_create.return_value = (MagicMock(), MagicMock(), MagicMock())
+        mock_create.return_value = (MagicMock(), MagicMock())
 
         MegatronLLMDeployable(
             megatron_checkpoint_filepath="model.ckpt",
@@ -519,7 +499,7 @@ def test_defaults_etp_1_and_sp_false(self, mock_create):
         """Defaults: expert_tensor_parallel_size=1, sequence_parallel=False."""
         from nemo_deploy.llm.megatronllm_deployable import MegatronLLMDeployable
 
-        mock_create.return_value = (MagicMock(), MagicMock(), MagicMock())
+        mock_create.return_value = (MagicMock(), MagicMock())
 
         MegatronLLMDeployable(megatron_checkpoint_filepath="model.ckpt")
 
diff --git a/tests/unit_tests/deploy/test_inference_base.py b/tests/unit_tests/deploy/test_inference_base.py
index d80c7dd2c..c35f4edea 100644
--- a/tests/unit_tests/deploy/test_inference_base.py
+++ b/tests/unit_tests/deploy/test_inference_base.py
@@ -19,10 +19,7 @@
 
 import pytest
 import torch
-from megatron.core.inference.engines.mcore_engine import MCoreEngine
-from megatron.core.inference.model_inference_wrappers.gpt.gpt_inference_wrapper import (
-    GPTInferenceWrapper,
-)
+from megatron.core.inference.apis import MegatronLLM
 from megatron.core.transformer.module import MegatronModule
 
 from nemo_deploy.llm.inference.inference_base import (
@@ -322,33 +319,31 @@ def test_setup_model_and_tokenizer_not_dist_ckpt(
         mock_check_dist.assert_called_once()
 
     def test_mcore_engine_with_cleanup(self):
-        # Create mocks for the engine and wrapper
-        mock_engine = MagicMock(spec=MCoreEngine)
-        mock_wrapper = MagicMock(spec=GPTInferenceWrapper)
+        # Create mock for the LLM engine
+        mock_llm = MagicMock(spec=MegatronLLM)
 
         # Create the wrapper
-        engine_wrapper = MCoreEngineWithCleanup(mock_engine, mock_wrapper, self.mock_tokenizer)
+        engine_wrapper = MCoreEngineWithCleanup(mock_llm, self.mock_tokenizer)
 
         # Test attribute delegation - mock the attribute access directly instead of using __getattr__
         # Define the attribute directly on the mock
-        mock_engine.some_attribute = "attribute_value"
+        mock_llm.some_attribute = "attribute_value"
         attribute_value = engine_wrapper.some_attribute
         self.assertEqual(attribute_value, "attribute_value")
 
         # Test method delegation - create a method on the mock
-        mock_engine.some_method = MagicMock(return_value="method_result")
+        mock_llm.some_method = MagicMock(return_value="method_result")
         result = engine_wrapper.some_method()
         self.assertEqual(result, "method_result")
-        mock_engine.some_method.assert_called_once()
+        mock_llm.some_method.assert_called_once()
 
     @patch("nemo_deploy.llm.inference.inference_base.cleanup_distributed")
     def test_mcore_engine_with_cleanup_del(self, mock_cleanup):
-        # Create mocks
-        mock_engine = MagicMock(spec=MCoreEngine)
-        mock_wrapper = MagicMock(spec=GPTInferenceWrapper)
+        # Create mock for the LLM engine
+        mock_llm = MagicMock(spec=MegatronLLM)
 
         # Create the wrapper
-        engine_wrapper = MCoreEngineWithCleanup(mock_engine, mock_wrapper, self.mock_tokenizer)
+        engine_wrapper = MCoreEngineWithCleanup(mock_llm, self.mock_tokenizer)
 
         # Call __del__
         engine_wrapper.__del__()
@@ -781,26 +776,20 @@ def test_setup_model_and_tokenizer_model_config_kwargs(
         self.assertEqual(self.model_config.hidden_size, 1024)
 
     @patch("nemo_deploy.llm.inference.inference_base.setup_model_and_tokenizer_for_inference")
-    @patch("nemo_deploy.llm.inference.inference_base.StaticInferenceContext")
-    @patch("nemo_deploy.llm.inference.inference_base.GPTInferenceWrapper")
-    @patch("nemo_deploy.llm.inference.inference_base.TextGenerationController")
-    @patch("nemo_deploy.llm.inference.inference_base.MCoreEngine")
+    @patch("nemo_deploy.llm.inference.inference_base.MegatronLLM")
     def test_create_mcore_engine_nemo_format(
         self,
-        mock_mcore_engine,
-        mock_text_ctrl,
-        mock_gpt_wrapper,
-        mock_static_ctx,
+        mock_megatron_llm,
         mock_setup,
     ):
         """Test create_mcore_engine with nemo model_format."""
         mock_model = MagicMock()
         mock_tokenizer = MagicMock()
         mock_setup.return_value = ([mock_model], mock_tokenizer)
-        mock_engine_instance = MagicMock()
-        mock_mcore_engine.return_value = mock_engine_instance
+        mock_llm_instance = MagicMock()
+        mock_megatron_llm.return_value = mock_llm_instance
 
-        engine, wrapper, tokenizer = create_mcore_engine(
+        engine, tokenizer = create_mcore_engine(
             path=self.mock_path,
             model_format="nemo",
             inference_max_seq_length=2048,
@@ -808,20 +797,14 @@ def test_create_mcore_engine_nemo_format(
         )
 
         mock_setup.assert_called_once()
-        mock_mcore_engine.assert_called_once()
+        mock_megatron_llm.assert_called_once()
         self.assertIsNotNone(engine)
 
     @patch("nemo_deploy.llm.inference.inference_base.setup_megatron_model_and_tokenizer_for_inference")
-    @patch("nemo_deploy.llm.inference.inference_base.StaticInferenceContext")
-    @patch("nemo_deploy.llm.inference.inference_base.GPTInferenceWrapper")
-    @patch("nemo_deploy.llm.inference.inference_base.TextGenerationController")
-    @patch("nemo_deploy.llm.inference.inference_base.MCoreEngine")
+    @patch("nemo_deploy.llm.inference.inference_base.MegatronLLM")
     def test_create_mcore_engine_megatron_format(
         self,
-        mock_mcore_engine,
-        mock_text_ctrl,
-        mock_gpt_wrapper,
-        mock_static_ctx,
+        mock_megatron_llm,
         mock_setup,
     ):
         """Test create_mcore_engine with megatron model_format."""
@@ -829,10 +812,10 @@ def test_create_mcore_engine_megatron_format(
         mock_tokenizer = MagicMock()
         mock_mlm_args = MagicMock()
         mock_setup.return_value = ([mock_model], mock_tokenizer, mock_mlm_args)
-        mock_engine_instance = MagicMock()
-        mock_mcore_engine.return_value = mock_engine_instance
+        mock_llm_instance = MagicMock()
+        mock_megatron_llm.return_value = mock_llm_instance
 
-        engine, wrapper, tokenizer = create_mcore_engine(
+        engine, tokenizer = create_mcore_engine(
             path=self.mock_path,
             model_format="megatron",
             inference_max_seq_length=2048,
@@ -840,7 +823,7 @@ def test_create_mcore_engine_megatron_format(
         )
 
         mock_setup.assert_called_once()
-        mock_mcore_engine.assert_called_once()
+        mock_megatron_llm.assert_called_once()
         self.assertIsNotNone(engine)
 
     @patch("nemo_deploy.llm.inference.inference_base.torch_distributed_init")
diff --git a/tests/unit_tests/deploy/test_megatron_multimodal_deployable.py b/tests/unit_tests/deploy/test_megatron_multimodal_deployable.py
index 666e4779c..81a6c760e 100644
--- a/tests/unit_tests/deploy/test_megatron_multimodal_deployable.py
+++ b/tests/unit_tests/deploy/test_megatron_multimodal_deployable.py
@@ -18,7 +18,7 @@
 import numpy as np
 import pytest
 import torch
-from megatron.core.inference.common_inference_params import CommonInferenceParams
+from megatron.core.inference.sampling_params import SamplingParams
 from PIL import Image
 
 from nemo_deploy.multimodal.megatron_multimodal_deployable import MegatronMultimodalDeployable
@@ -157,7 +157,7 @@ def test_generate_method(self, deployable, sample_image):
         """Test the generate method."""
         prompts = ["Test prompt 1", "Test prompt 2"]
         images = [sample_image, sample_image]
-        inference_params = CommonInferenceParams(temperature=0.7, top_k=10, top_p=0.9, num_tokens_to_generate=100)
+        inference_params = SamplingParams(temperature=0.7, top_k=10, top_p=0.9, num_tokens_to_generate=100)
 
         with patch("nemo_deploy.multimodal.megatron_multimodal_deployable.generate") as mock_generate:
             with patch.object(deployable, "apply_chat_template", side_effect=lambda x: x):
@@ -282,8 +282,8 @@ def test_infer_fn(self, deployable, sample_image_base64, sample_image):
                 assert call_args[0][0] == prompts
                 # Images should be converted from base64
                 assert len(call_args[0][1]) == 2
-                # Check that inference_params is a CommonInferenceParams object (3rd positional arg)
-                assert isinstance(call_args[0][2], CommonInferenceParams)
+                # Check that inference_params is a SamplingParams object (3rd positional arg)
+                assert isinstance(call_args[0][2], SamplingParams)
                 assert call_args[0][2].temperature == 0.8
                 assert call_args[0][2].top_k == 20
                 assert call_args[0][2].top_p == 0.95
@@ -318,8 +318,8 @@ def test_infer_fn_default_params(self, deployable, sample_image_base64, sample_i
                 assert call_args[0][0] == prompts
                 # Images should be converted from base64
                 assert len(call_args[0][1]) == 1
-                # Check that inference_params is a CommonInferenceParams object (3rd positional arg)
-                assert isinstance(call_args[0][2], CommonInferenceParams)
+                # Check that inference_params is a SamplingParams object (3rd positional arg)
+                assert isinstance(call_args[0][2], SamplingParams)
                 assert call_args[0][2].temperature == 1.0
                 assert call_args[0][2].top_k == 1
                 assert call_args[0][2].top_p == 0.0
@@ -357,7 +357,7 @@ def test_infer_fn_with_temperature_zero(self, deployable):
                 call_args = mock_generate.call_args
 
                 # Check that inference_params has greedy sampling parameters
-                assert isinstance(call_args[0][2], CommonInferenceParams)
+                assert isinstance(call_args[0][2], SamplingParams)
                 assert call_args[0][2].temperature == 0.0  # Kept as 0.0
                 assert call_args[0][2].top_k == 1  # Overridden for greedy sampling
                 assert call_args[0][2].top_p == 0.0  # Overridden for greedy sampling
diff --git a/tests/unit_tests/deploy/test_megatronllm_deployable.py b/tests/unit_tests/deploy/test_megatronllm_deployable.py
index af6902770..0e1ef4d4f 100644
--- a/tests/unit_tests/deploy/test_megatronllm_deployable.py
+++ b/tests/unit_tests/deploy/test_megatronllm_deployable.py
@@ -16,7 +16,7 @@
 
 import numpy as np
 import pytest
-from megatron.core.inference.common_inference_params import CommonInferenceParams
+from megatron.core.inference.sampling_params import SamplingParams
 
 from nemo_deploy.llm.megatronllm_deployable import MegatronLLMDeployable, dict_to_str
 from nemo_export_deploy_common.import_utils import UnavailableError
@@ -121,7 +121,7 @@ def test_generate_with_cuda_graphs_empty_prompts(deployable):
     deployable.enable_cuda_graphs = True
     deployable.max_batch_size = 4
     prompts = []
-    inference_params = CommonInferenceParams()
+    inference_params = SamplingParams()
 
     with patch.object(deployable.mcore_engine, "generate") as mock_generate:
         mock_generate.return_value = ["", "", "", ""]
@@ -197,13 +197,13 @@ def test_generate_other_ranks_disables_materialize_when_log_probs(deployable):
             [1.0, 1, 0.0, 256, True, None],  # log_probs=True
         ]
 
-        deployable.mcore_engine.dynamic_engine.materialize_only_last_token_logits = True
-        deployable.mcore_engine.dynamic_engine.context.config.materialize_only_last_token_logits = True
+        deployable.mcore_engine.engine.materialize_only_last_token_logits = True
+        deployable.mcore_engine.engine.context.config.materialize_only_last_token_logits = True
 
         deployable.generate_other_ranks()
 
-        assert deployable.mcore_engine.dynamic_engine.materialize_only_last_token_logits is False
-        assert deployable.mcore_engine.dynamic_engine.context.config.materialize_only_last_token_logits is False
+        assert deployable.mcore_engine.engine.materialize_only_last_token_logits is False
+        assert deployable.mcore_engine.engine.context.config.materialize_only_last_token_logits is False
 
 
 @pytest.mark.run_only_on("GPU")
@@ -387,7 +387,7 @@ def test_generate_without_cuda_graphs(deployable):
     deployable.enable_cuda_graphs = False
 
     prompts = ["Hello", "World"]
-    inference_params = CommonInferenceParams(
+    inference_params = SamplingParams(
         temperature=1.0,
         top_k=1,
         top_p=0.0,
@@ -403,7 +403,7 @@ def test_generate_without_cuda_graphs(deployable):
 
         results = deployable.generate(prompts, inference_params)
         assert len(results) == 2
-        mock_generate.assert_called_once_with(prompts=prompts, add_BOS=False, common_inference_params=inference_params)
+        mock_generate.assert_called_once_with(prompts=prompts, sampling_params=inference_params)
 
 
 @pytest.mark.run_only_on("GPU")
@@ -414,7 +414,7 @@ def test_generate_with_cuda_graphs(deployable):
     deployable.max_batch_size = 4
 
     prompts = ["Hello", "World"]
-    inference_params = CommonInferenceParams(
+    inference_params = SamplingParams(
         temperature=1.0,
         top_k=1,
         top_p=0.0,
@@ -446,8 +446,7 @@ def test_generate_with_cuda_graphs(deployable):
         called_args = mock_generate.call_args[1]
         assert len(called_args["prompts"]) == 4  # Should pad to max_batch_size
         assert called_args["prompts"][:2] == prompts  # Original prompts should be first
-        assert called_args["add_BOS"] is False
-        assert called_args["common_inference_params"] == inference_params
+        assert called_args["sampling_params"] == inference_params
 
 
 @pytest.mark.run_only_on("GPU")
@@ -607,8 +606,8 @@ def test_init_with_megatron_valid_types(model_type):
         patch("nemo_deploy.llm.megatronllm_deployable.HAVE_TRITON", True),
         patch("nemo_deploy.llm.megatronllm_deployable.create_mcore_engine") as mock_create,
     ):
-        mock_engine, mock_model, mock_tokenizer = MagicMock(), MagicMock(), MagicMock()
-        mock_create.return_value = (mock_engine, mock_model, mock_tokenizer)
+        mock_engine, mock_tokenizer = MagicMock(), MagicMock()
+        mock_create.return_value = (mock_engine, mock_tokenizer)
 
         deployable = MegatronLLMDeployable(
             megatron_checkpoint_filepath="bar.ckpt",
@@ -716,7 +715,7 @@ def test_infer_fn_basic(deployable):
         call_args = mock_generate.call_args[0]
         assert call_args[0] == prompts
 
-        # Verify CommonInferenceParams
+        # Verify SamplingParams
         inference_params = mock_generate.call_args[0][1]
         assert inference_params.temperature == 1.0
         assert inference_params.top_k == 1
@@ -1076,13 +1075,13 @@ def test_infer_fn_disables_materialize_only_last_token_logits_when_log_probs(dep
         mock_tensor_instance.cpu.return_value.detach.return_value.numpy.return_value = np.array([0.1])
         mock_tensor.return_value = mock_tensor_instance
 
-        deployable.mcore_engine.dynamic_engine.materialize_only_last_token_logits = True
-        deployable.mcore_engine.dynamic_engine.context.config.materialize_only_last_token_logits = True
+        deployable.mcore_engine.engine.materialize_only_last_token_logits = True
+        deployable.mcore_engine.engine.context.config.materialize_only_last_token_logits = True
 
         deployable._infer_fn(prompts=["Hello"], log_probs=True)
 
-        assert deployable.mcore_engine.dynamic_engine.materialize_only_last_token_logits is False
-        assert deployable.mcore_engine.dynamic_engine.context.config.materialize_only_last_token_logits is False
+        assert deployable.mcore_engine.engine.materialize_only_last_token_logits is False
+        assert deployable.mcore_engine.engine.context.config.materialize_only_last_token_logits is False
 
 
 @pytest.mark.run_only_on("GPU")
@@ -1100,13 +1099,13 @@ def test_infer_fn_disables_materialize_only_last_token_logits_when_top_logprobs(
         mock_remove_eos.return_value = ["text"]
         mock_dict_to_str.return_value = '{"tok": 0.1}'
 
-        deployable.mcore_engine.dynamic_engine.materialize_only_last_token_logits = True
-        deployable.mcore_engine.dynamic_engine.context.config.materialize_only_last_token_logits = True
+        deployable.mcore_engine.engine.materialize_only_last_token_logits = True
+        deployable.mcore_engine.engine.context.config.materialize_only_last_token_logits = True
 
         deployable._infer_fn(prompts=["Hello"], top_logprobs=5)
 
-        assert deployable.mcore_engine.dynamic_engine.materialize_only_last_token_logits is False
-        assert deployable.mcore_engine.dynamic_engine.context.config.materialize_only_last_token_logits is False
+        assert deployable.mcore_engine.engine.materialize_only_last_token_logits is False
+        assert deployable.mcore_engine.engine.context.config.materialize_only_last_token_logits is False
 
 
 @pytest.mark.run_only_on("GPU")
@@ -1121,10 +1120,10 @@ def test_infer_fn_keeps_materialize_only_last_token_logits_when_no_logprobs(depl
         mock_generate.return_value = [mock_result]
         mock_remove_eos.return_value = ["text"]
 
-        deployable.mcore_engine.dynamic_engine.materialize_only_last_token_logits = True
-        deployable.mcore_engine.dynamic_engine.context.config.materialize_only_last_token_logits = True
+        deployable.mcore_engine.engine.materialize_only_last_token_logits = True
+        deployable.mcore_engine.engine.context.config.materialize_only_last_token_logits = True
 
         deployable._infer_fn(prompts=["Hello"], log_probs=False, top_logprobs=0)
 
-        assert deployable.mcore_engine.dynamic_engine.materialize_only_last_token_logits is True
-        assert deployable.mcore_engine.dynamic_engine.context.config.materialize_only_last_token_logits is True
+        assert deployable.mcore_engine.engine.materialize_only_last_token_logits is True
+        assert deployable.mcore_engine.engine.context.config.materialize_only_last_token_logits is True

From 8704af9d5d65023ecd145f73b94b6e2ca0be5ebe Mon Sep 17 00:00:00 2001
From: Onur Yilmaz <oyilmaz@nvidia.com>
Date: Wed, 1 Jul 2026 06:36:43 -0400
Subject: [PATCH 2/5] Update dependencies

Signed-off-by: Onur Yilmaz <oyilmaz@nvidia.com>
---
 pyproject.toml |   4 +-
 uv.lock        | 141 +++++++++++++++++++++++++------------------------
 2 files changed, 75 insertions(+), 70 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 319a588f5..ca57f5210 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -110,7 +110,7 @@ vllm = [
     { index = "pytorch-cu130", marker = "python_version < '3.9' and platform_machine == 'x86_64'" },
     { index = "pypi", marker = "platform_machine == 'aarch64'" },
 ]
-megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "859a18d001eca4423ad79e7771632d839f50fa77" }
+megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "97e8dba4060ae74a110b9da6255eb0a66b016b2d" }
 nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", rev = "b2bb3d728a18795807d9f76c535e005a609a1b01" }
 # nemo-toolkit = { git = "https://github.com/NVIDIA/NeMo.git", rev = "main" }
 
@@ -124,6 +124,7 @@ no-build-isolation-package = [
     "causal-conv1d",
     "nv-grouped-gemm",
 ]
+
 # Always apply the build group since dependencies like TE/mcore/nemo-run require build dependencies
 # and this lets us assume they are implicitly installed with a simply `uv sync`. Ideally, we'd
 # avoid including these in the default dependency set, but for now it's required.
@@ -134,6 +135,7 @@ default-groups = ["linting", "build", "test"]
 #  --link-mode=symlink (fastest option when uv cache and venv on different file-system; caveat: venv is brittle since it depends on the environment/container)
 link-mode = "copy"
 conflicts = [[{ extra = "vllm" }, { extra = "trt-onnx" }]]
+extra-build-dependencies = { fast-hadamard-transform = ["setuptools", "torch"] }
 override-dependencies = [
     "cachetools>=5.0.0",
     "torch; sys_platform == 'never'",
diff --git a/uv.lock b/uv.lock
index 5b03673ec..c8c804f13 100644
--- a/uv.lock
+++ b/uv.lock
@@ -678,7 +678,8 @@ dependencies = [
     { name = "rich" },
     { name = "semantic-version" },
     { name = "sentry-sdk" },
-    { name = "setuptools" },
+    { name = "setuptools", version = "79.0.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "simplejson" },
     { name = "urllib3" },
     { name = "wrapt" },
@@ -1207,6 +1208,16 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d6/1f/e99e23ee01847147fa194e8d41cfcf2535a2dbfcb51414c541cadb15c5d7/fabric-3.2.2-py3-none-any.whl", hash = "sha256:91c47c0be68b14936c88b34da8a1f55e5710fd28397dac5d4ff2e21558113a6f", size = 59417, upload-time = "2023-08-31T01:42:03.917Z" },
 ]
 
+[[package]]
+name = "fast-hadamard-transform"
+version = "1.0.4.post1"
+source = { git = "https://github.com/Dao-AILab/fast-hadamard-transform.git?rev=f134af63deb2df17e1171a9ec1ea4a7d8604d5ca#f134af63deb2df17e1171a9ec1ea4a7d8604d5ca" }
+dependencies = [
+    { name = "ninja" },
+    { name = "packaging" },
+    { name = "torch", marker = "sys_platform == 'never' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+]
+
 [[package]]
 name = "fastapi"
 version = "0.121.3"
@@ -1308,11 +1319,11 @@ wheels = [
 
 [[package]]
 name = "filelock"
-version = "3.19.1"
+version = "3.29.4"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/40/bb/0ab3e58d22305b6f5440629d20683af28959bf793d98d11950e305c1c326/filelock-3.19.1.tar.gz", hash = "sha256:66eda1888b0171c998b35be2bcc0f6d75c388a7ce20c3f3f37aa8e96c2dddf58", size = 17687, upload-time = "2025-08-14T16:56:03.016Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/e6/dc/be6cbe99670cd6e4ad387123647cb08e0c32975e223f82551e914c5568a6/filelock-3.29.4.tar.gz", hash = "sha256:10cdb3656fc44541cdf30652a93fb10ec6b05325620eb316bd26893e4201538a", size = 63028, upload-time = "2026-06-13T16:12:00.744Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/42/14/42b2651a2f46b022ccd948bca9f2d5af0fd8929c4eec235b8d6d844fbe67/filelock-3.19.1-py3-none-any.whl", hash = "sha256:d38e30481def20772f5baf097c122c3babc4fcdb7e14e57049eb9d88c6dc017d", size = 15988, upload-time = "2025-08-14T16:56:01.633Z" },
+    { url = "https://files.pythonhosted.org/packages/13/37/a065dc3bd6e49423a6532c642ca7378d3f467b1ef44c2800c937af7f9739/filelock-3.29.4-py3-none-any.whl", hash = "sha256:dac1648087d5115554850d113e7dd8c83ab2d38e3435dde2d4f163847e57b767", size = 42757, upload-time = "2026-06-13T16:11:59.582Z" },
 ]
 
 [[package]]
@@ -1705,7 +1716,8 @@ source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "grpcio" },
     { name = "protobuf" },
-    { name = "setuptools" },
+    { name = "setuptools", version = "79.0.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/8b/d1/cbefe328653f746fd319c4377836a25ba64226e41c6a1d7d5cdbc87a459f/grpcio_tools-1.78.0.tar.gz", hash = "sha256:4b0dd86560274316e155d925158276f8564508193088bc43e20d3f5dff956b2b", size = 5393026, upload-time = "2026-02-06T09:59:59.53Z" }
 wheels = [
@@ -1949,33 +1961,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/66/13/530802bc391c95be6fe9f96e9aa427d94067e7c0b7da7a9092344dc44c4b/ijson-3.4.0-cp312-cp312-win_amd64.whl", hash = "sha256:71523f2b64cb856a820223e94d23e88369f193017ecc789bb4de198cc9d349eb", size = 54081, upload-time = "2025-05-08T02:36:07.099Z" },
 ]
 
-[[package]]
-name = "imageio"
-version = "2.37.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "numpy" },
-    { name = "pillow" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/0c/47/57e897fb7094afb2d26e8b2e4af9a45c7cf1a405acdeeca001fdf2c98501/imageio-2.37.0.tar.gz", hash = "sha256:71b57b3669666272c818497aebba2b4c5f20d5b37c81720e5e1a56d59c492996", size = 389963, upload-time = "2025-01-20T02:42:37.089Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/cb/bd/b394387b598ed84d8d0fa90611a90bee0adc2021820ad5729f7ced74a8e2/imageio-2.37.0-py3-none-any.whl", hash = "sha256:11efa15b87bc7871b61590326b2d635439acc321cf7f8ce996f812543ce10eed", size = 315796, upload-time = "2025-01-20T02:42:34.931Z" },
-]
-
-[[package]]
-name = "imageio-ffmpeg"
-version = "0.6.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" },
-    { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" },
-    { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" },
-    { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" },
-    { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" },
-    { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" },
-]
-
 [[package]]
 name = "imagesize"
 version = "1.4.1"
@@ -2325,7 +2310,8 @@ version = "0.15.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "packaging" },
-    { name = "setuptools" },
+    { name = "setuptools", version = "79.0.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "typing-extensions" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/b8/39/6fc58ca81492db047149b4b8fd385aa1bfb8c28cd7cacb0c7eb0c44d842f/lightning_utilities-0.15.2.tar.gz", hash = "sha256:cdf12f530214a63dacefd713f180d1ecf5d165338101617b4742e8f22c032e24", size = 31090, upload-time = "2025-08-06T13:57:39.242Z" }
@@ -2533,8 +2519,8 @@ wheels = [
 
 [[package]]
 name = "megatron-bridge"
-version = "0.5.0+859a18d0"
-source = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git?rev=859a18d001eca4423ad79e7771632d839f50fa77#859a18d001eca4423ad79e7771632d839f50fa77" }
+version = "0.6.0+97e8dba4"
+source = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git?rev=97e8dba4060ae74a110b9da6255eb0a66b016b2d#97e8dba4060ae74a110b9da6255eb0a66b016b2d" }
 dependencies = [
     { name = "accelerate" },
     { name = "comet-ml" },
@@ -2542,9 +2528,9 @@ dependencies = [
     { name = "diffusers" },
     { name = "einops" },
     { name = "flash-linear-attention" },
+    { name = "flashinfer-cubin" },
+    { name = "flashinfer-python" },
     { name = "hydra-core" },
-    { name = "imageio" },
-    { name = "imageio-ffmpeg" },
     { name = "megatron-core", extra = ["dev", "mlm"] },
     { name = "mistral-common" },
     { name = "mlflow" },
@@ -2568,8 +2554,8 @@ dependencies = [
 
 [[package]]
 name = "megatron-core"
-version = "0.18.0+5c7968afa"
-source = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git?subdirectory=3rdparty%2FMegatron-LM&rev=859a18d001eca4423ad79e7771632d839f50fa77#859a18d001eca4423ad79e7771632d839f50fa77" }
+version = "0.19.0+b9f7fb68a"
+source = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git?subdirectory=3rdparty%2FMegatron-LM&rev=97e8dba4060ae74a110b9da6255eb0a66b016b2d#97e8dba4060ae74a110b9da6255eb0a66b016b2d" }
 dependencies = [
     { name = "numpy" },
     { name = "packaging" },
@@ -2582,6 +2568,7 @@ dev = [
     { name = "datasets" },
     { name = "einops" },
     { name = "emerging-optimizers" },
+    { name = "fast-hadamard-transform" },
     { name = "fastapi" },
     { name = "flash-linear-attention" },
     { name = "flashinfer-python" },
@@ -2598,10 +2585,12 @@ dev = [
     { name = "tensorstore" },
     { name = "tqdm" },
     { name = "wget" },
+    { name = "zstandard" },
 ]
 mlm = [
     { name = "accelerate" },
     { name = "flask-restful" },
+    { name = "omegaconf" },
     { name = "sentencepiece" },
     { name = "tiktoken" },
     { name = "transformers" },
@@ -2769,7 +2758,8 @@ dependencies = [
     { name = "httpx" },
     { name = "jmespath" },
     { name = "pydantic" },
-    { name = "setuptools" },
+    { name = "setuptools", version = "79.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'darwin' and extra == 'extra-18-nemo-export-deploy-vllm') or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'darwin' and extra == 'extra-18-nemo-export-deploy-vllm') or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "starlette" },
     { name = "supervisor" },
 ]
@@ -2822,7 +2812,7 @@ wheels = [
 
 [[package]]
 name = "multi-storage-client"
-version = "0.30.0"
+version = "0.50.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "filelock" },
@@ -2835,14 +2825,14 @@ dependencies = [
     { name = "python-dateutil" },
     { name = "pyyaml" },
     { name = "tqdm" },
+    { name = "tzdata" },
     { name = "wcmatch" },
     { name = "xattr" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/ad/13/5af5e56066c8a4e0db01e05a1dc1209143d46a69b8d5119d75aaca5470e1/multi_storage_client-0.30.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:48a582dae6510a1e8511ac36a98a6c929e60e88177abcdd00dda637d4728a478", size = 5265831, upload-time = "2025-09-26T18:01:21.464Z" },
-    { url = "https://files.pythonhosted.org/packages/c6/fa/9e09e93c06955d966101882eb9f9998bde2738a77bf627c32df385f06d80/multi_storage_client-0.30.0-cp312-cp312-macosx_11_0_x86_64.whl", hash = "sha256:0fada9795b61d680ff7532a76d270a0accda134c3da365915721f808c1bbbc41", size = 5389337, upload-time = "2025-09-26T17:58:30.981Z" },
-    { url = "https://files.pythonhosted.org/packages/4f/59/edf848db15176cdd00204356627bf5aefded67336a57945cf70733356ba2/multi_storage_client-0.30.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cae72686e96f75a1d03e74da46ca2f95ba6d086a3091529bdecd42944d91884a", size = 3168948, upload-time = "2025-09-26T17:58:54.478Z" },
-    { url = "https://files.pythonhosted.org/packages/92/d6/51f54871c05fe6b493623f51c29beb7e92ad96dbc2dbe0577136c6ea8dd2/multi_storage_client-0.30.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5bbede6b73ee48c24d68dd5700fc7637ce9aaf479f74ccb3d4a4a5a4f7a1b6d4", size = 3342293, upload-time = "2025-09-26T17:57:44.292Z" },
+    { url = "https://files.pythonhosted.org/packages/30/cd/5e37c93d73f5df512f2e452a9381e5da3d818b385c6f6a612e1e65d48e83/multi_storage_client-0.50.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:a164b7d56d238323f4a194834ff69969b7c036bb87442e0bdc39ed5754d6c786", size = 10109032, upload-time = "2026-06-05T22:00:01.115Z" },
+    { url = "https://files.pythonhosted.org/packages/da/cc/f68f7b9a70d43e2e7616c68893153ec7429bd5fd28b3c077dd3ac9dedf96/multi_storage_client-0.50.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:97c3b96956ce3296ce1c77c1cbe7680f310780d652593c8dfd35f2556b686e28", size = 6340158, upload-time = "2026-06-05T22:00:28.497Z" },
+    { url = "https://files.pythonhosted.org/packages/9d/58/73978d412abdf82cacf8346c9cefaa4310b354066ac8cfc15b169805481a/multi_storage_client-0.50.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4b5691e23e16d85fc29a7467d481e4a987d6203cda99630a90d5b033419b7f80", size = 6539022, upload-time = "2026-06-05T22:00:48.237Z" },
 ]
 
 [[package]]
@@ -2972,7 +2962,8 @@ build = [
     { name = "cython" },
     { name = "ninja" },
     { name = "pybind11" },
-    { name = "setuptools" },
+    { name = "setuptools", version = "79.0.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "torch", marker = "sys_platform == 'never' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
 ]
 docs = [
@@ -3009,7 +3000,7 @@ requires-dist = [
     { name = "hydra-core", specifier = ">1.3,<=1.3.2" },
     { name = "ijson" },
     { name = "lightning", specifier = "<2.5.0" },
-    { name = "megatron-bridge", git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git?rev=859a18d001eca4423ad79e7771632d839f50fa77" },
+    { name = "megatron-bridge", git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git?rev=97e8dba4060ae74a110b9da6255eb0a66b016b2d" },
     { name = "megatron-core" },
     { name = "nvidia-modelopt", extras = ["torch"], marker = "sys_platform != 'darwin'" },
     { name = "nvidia-pytriton", marker = "sys_platform != 'darwin'" },
@@ -3394,25 +3385,27 @@ wheels = [
 
 [[package]]
 name = "nvidia-modelopt"
-version = "0.37.0"
+version = "0.44.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "ninja", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "numpy", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "nvidia-ml-py", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "omegaconf", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "packaging", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "pulp", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "pydantic", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "pyyaml", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "regex", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "rich", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "safetensors", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "scipy", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "torch", marker = "sys_platform == 'never' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
-    { name = "torchprofile", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "tqdm", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/27/a8/85034a33753a56ef120b931dda1183d98efef1010e379c54d3c3214a5048/nvidia_modelopt-0.37.0-py3-none-any.whl", hash = "sha256:3490b6d6aea3541aa5d475d81230fee627e2c16ff47bbab1cba4b80a1eb119a2", size = 831271, upload-time = "2025-10-08T18:37:03.951Z" },
+    { url = "https://files.pythonhosted.org/packages/83/ab/7e12dd238638624cb9d48904e9205abe16a5b26bdd5f9b91e3357821cf90/nvidia_modelopt-0.44.0-py3-none-any.whl", hash = "sha256:9b54a853dfda161db97a0dfce4d7c24269d1d19966b5e2026094186af897f74d", size = 1604658, upload-time = "2026-05-13T20:47:04.007Z" },
 ]
 
 [[package]]
@@ -5015,11 +5008,33 @@ wheels = [
 name = "setuptools"
 version = "79.0.1"
 source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "platform_machine == 'aarch64' and sys_platform == 'darwin' and extra != 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm'",
+    "platform_machine != 'aarch64' and sys_platform == 'darwin' and extra != 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm'",
+    "sys_platform == 'darwin' and extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra != 'extra-18-nemo-export-deploy-vllm'",
+    "sys_platform == 'darwin' and extra != 'extra-18-nemo-export-deploy-trt-onnx' and extra != 'extra-18-nemo-export-deploy-vllm'",
+]
 sdist = { url = "https://files.pythonhosted.org/packages/bb/71/b6365e6325b3290e14957b2c3a804a529968c77a049b2ed40c095f749707/setuptools-79.0.1.tar.gz", hash = "sha256:128ce7b8f33c3079fd1b067ecbb4051a66e8526e7b65f6cec075dfc650ddfa88", size = 1367909, upload-time = "2025-04-23T22:20:59.241Z" }
 wheels = [
     { url = "https://files.pythonhosted.org/packages/0d/6d/b4752b044bf94cb802d88a888dc7d288baaf77d7910b7dedda74b5ceea0c/setuptools-79.0.1-py3-none-any.whl", hash = "sha256:e147c0549f27767ba362f9da434eab9c5dc0045d5304feb602a0af001089fc51", size = 1256281, upload-time = "2025-04-23T22:20:56.768Z" },
 ]
 
+[[package]]
+name = "setuptools"
+version = "80.10.2"
+source = { registry = "https://pypi.org/simple" }
+resolution-markers = [
+    "platform_machine == 'aarch64' and sys_platform == 'linux' and extra != 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm'",
+    "platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'linux' and extra != 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm'",
+    "platform_machine != 'aarch64' and sys_platform != 'darwin' and extra != 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm'",
+    "sys_platform != 'darwin' and extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra != 'extra-18-nemo-export-deploy-vllm'",
+    "sys_platform != 'darwin' and extra != 'extra-18-nemo-export-deploy-trt-onnx' and extra != 'extra-18-nemo-export-deploy-vllm'",
+]
+sdist = { url = "https://files.pythonhosted.org/packages/76/95/faf61eb8363f26aa7e1d762267a8d602a1b26d4f3a1e758e92cb3cb8b054/setuptools-80.10.2.tar.gz", hash = "sha256:8b0e9d10c784bf7d262c4e5ec5d4ec94127ce206e8738f29a437945fbc219b70", size = 1200343, upload-time = "2026-01-25T22:38:17.252Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/94/b8/f1f62a5e3c0ad2ff1d189590bfa4c46b4f3b6e49cef6f26c6ee4e575394d/setuptools-80.10.2-py3-none-any.whl", hash = "sha256:95b30ddfb717250edb492926c92b5221f7ef3fbcc2b07579bcd4a27da21d0173", size = 1064234, upload-time = "2026-01-25T22:38:15.216Z" },
+]
+
 [[package]]
 name = "sh"
 version = "2.2.2"
@@ -5391,7 +5406,8 @@ dependencies = [
     { name = "packaging" },
     { name = "pillow" },
     { name = "protobuf" },
-    { name = "setuptools" },
+    { name = "setuptools", version = "79.0.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "tensorboard-data-server" },
     { name = "werkzeug" },
 ]
@@ -5498,7 +5514,7 @@ dependencies = [
     { name = "ml-dtypes" },
     { name = "numpy" },
     { name = "psutil" },
-    { name = "setuptools", marker = "sys_platform == 'darwin'" },
+    { name = "setuptools", version = "79.0.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" },
     { name = "torch", marker = "sys_platform == 'never'" },
     { name = "torch-c-dlpack-ext" },
     { name = "tqdm" },
@@ -5577,7 +5593,7 @@ dependencies = [
     { name = "nvidia-nvjitlink", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform == 'linux' and extra == 'extra-18-nemo-export-deploy-trt-onnx') or (sys_platform == 'linux' and extra != 'extra-18-nemo-export-deploy-vllm') or (sys_platform != 'linux' and extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "nvidia-nvshmem-cu13", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform == 'linux' and extra == 'extra-18-nemo-export-deploy-trt-onnx') or (sys_platform == 'linux' and extra != 'extra-18-nemo-export-deploy-vllm') or (sys_platform != 'linux' and extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "nvidia-nvtx", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform == 'linux' and extra == 'extra-18-nemo-export-deploy-trt-onnx') or (sys_platform == 'linux' and extra != 'extra-18-nemo-export-deploy-vllm') or (sys_platform != 'linux' and extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
-    { name = "setuptools", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux') or (sys_platform == 'linux' and extra == 'extra-18-nemo-export-deploy-trt-onnx') or (sys_platform == 'linux' and extra != 'extra-18-nemo-export-deploy-vllm') or (sys_platform == 'darwin' and extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux') or (sys_platform == 'linux' and extra == 'extra-18-nemo-export-deploy-trt-onnx') or (sys_platform == 'linux' and extra != 'extra-18-nemo-export-deploy-vllm') or (sys_platform == 'darwin' and extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "sympy", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux') or (sys_platform == 'linux' and extra == 'extra-18-nemo-export-deploy-trt-onnx') or (sys_platform == 'linux' and extra != 'extra-18-nemo-export-deploy-vllm') or (sys_platform == 'darwin' and extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "triton", marker = "sys_platform == 'never' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "typing-extensions", marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux') or (sys_platform == 'linux' and extra == 'extra-18-nemo-export-deploy-trt-onnx') or (sys_platform == 'linux' and extra != 'extra-18-nemo-export-deploy-vllm') or (sys_platform == 'darwin' and extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
@@ -5629,20 +5645,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/02/21/aa0f434434c48490f91b65962b1ce863fdcce63febc166ca9fe9d706c2b6/torchmetrics-1.8.2-py3-none-any.whl", hash = "sha256:08382fd96b923e39e904c4d570f3d49e2cc71ccabd2a94e0f895d1f0dac86242", size = 983161, upload-time = "2025-09-03T14:00:51.921Z" },
 ]
 
-[[package]]
-name = "torchprofile"
-version = "0.0.4"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "numpy", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
-    { name = "torch", marker = "sys_platform == 'never' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
-    { name = "torchvision", marker = "sys_platform == 'never' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/6f/36/574c0c46e818533b78b3c09505211162918188325ab4165ef11a3f295755/torchprofile-0.0.4.tar.gz", hash = "sha256:96b6da17d752a06b02977e078aea95614893b31d4117dd5dcd081f30ce65611b", size = 4557, upload-time = "2021-06-22T04:58:03.592Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/62/15/71ad4ed163b03cba1315f1d96e0bc8e39d5a97f92974ffa610a729b273ab/torchprofile-0.0.4-py3-none-any.whl", hash = "sha256:7151fe88dc770f0eeec241244a4c7feaec2c5e8c7852386bc2d6a8d7dde7384d", size = 7694, upload-time = "2021-06-22T04:58:02.485Z" },
-]
-
 [[package]]
 name = "torchvision"
 version = "0.24.0"
@@ -5997,7 +5999,8 @@ dependencies = [
     { name = "requests" },
     { name = "sentencepiece" },
     { name = "setproctitle" },
-    { name = "setuptools" },
+    { name = "setuptools", version = "79.0.1", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform == 'darwin' and extra == 'extra-18-nemo-export-deploy-vllm') or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" }, marker = "(sys_platform != 'darwin' and extra == 'extra-18-nemo-export-deploy-vllm') or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
     { name = "six" },
     { name = "tiktoken" },
     { name = "tilelang" },
@@ -6342,7 +6345,7 @@ name = "zope-event"
 version = "6.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "setuptools", marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/c2/d8/9c8b0c6bb1db09725395618f68d3b8a08089fca0aed28437500caaf713ee/zope_event-6.0.tar.gz", hash = "sha256:0ebac894fa7c5f8b7a89141c272133d8c1de6ddc75ea4b1f327f00d1f890df92", size = 18731, upload-time = "2025-09-12T07:10:13.551Z" }
 wheels = [

From e67b52ccae1917c53f69a022ac766edee03abf68 Mon Sep 17 00:00:00 2001
From: Onur Yilmaz <oyilmaz@nvidia.com>
Date: Wed, 1 Jul 2026 06:51:16 -0400
Subject: [PATCH 3/5] Fix depenedency issue

Signed-off-by: Onur Yilmaz <oyilmaz@nvidia.com>
---
 pyproject.toml |   7 +-
 uv.lock        | 271 ++-----------------------------------------------
 2 files changed, 13 insertions(+), 265 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index ca57f5210..c0d5666d4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -135,7 +135,6 @@ default-groups = ["linting", "build", "test"]
 #  --link-mode=symlink (fastest option when uv cache and venv on different file-system; caveat: venv is brittle since it depends on the environment/container)
 link-mode = "copy"
 conflicts = [[{ extra = "vllm" }, { extra = "trt-onnx" }]]
-extra-build-dependencies = { fast-hadamard-transform = ["setuptools", "torch"] }
 override-dependencies = [
     "cachetools>=5.0.0",
     "torch; sys_platform == 'never'",
@@ -151,7 +150,7 @@ override-dependencies = [
     "datasets>=3.3.0",
     "flash-linear-attention>=0.3.0,<0.4.dev0",
     "patchelf; sys_platform=='never'",
-    "nvidia-resiliency-ext>=0.3.0,<0.6.0",
+    "nvidia-resiliency-ext>=0.3.0,<=0.6.0",
     "transformer-engine[pytorch,core_cu13]>=2.14.0a0,<2.17.0; sys_platform != 'darwin'",
     "transformer-engine-cu13>=2.14.0a0,<2.17.0; sys_platform != 'darwin'",
     "transformer-engine-cu12; sys_platform == 'never'",
@@ -164,6 +163,10 @@ override-dependencies = [
 ]
 prerelease = "allow"
 
+[[tool.uv.dependency-metadata]]
+name = "nvidia-resiliency-ext"
+version = "0.6.0"
+
 [[tool.uv.index]]
 name = "pypi"
 url = "https://pypi.org/simple"
diff --git a/uv.lock b/uv.lock
index c8c804f13..2ee58114f 100644
--- a/uv.lock
+++ b/uv.lock
@@ -48,6 +48,10 @@ overrides = [
     { name = "urllib3", specifier = ">1.27.0" },
 ]
 
+[[manifest.dependency-metadata]]
+name = "nvidia-resiliency-ext"
+version = "0.6.0"
+
 [[package]]
 name = "absl-py"
 version = "2.3.1"
@@ -981,15 +985,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/6c/be/e15b5b866da452e62635a7b27513f31cb581fa2ea9cc9b768b535d62a955/decord-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02665d7c4f1193a330205a791bc128f7e108eb6ae5b67144437a02f700943bad", size = 24733380, upload-time = "2021-06-14T21:30:57.766Z" },
 ]
 
-[[package]]
-name = "defusedxml"
-version = "0.8.0rc2"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/5e/3b/b8849dcc3f96913924137dc4ea041d74aa513a3c5dda83d8366491290c74/defusedxml-0.8.0rc2.tar.gz", hash = "sha256:138c7d540a78775182206c7c97fe65b246a2f40b29471e1a2f1b0da76e7a3942", size = 52575, upload-time = "2023-09-29T08:01:27.517Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/5d/c7/6b4ad89ca6f7732ff97ce5e9caa6fe739600d26c5d53c20d0bf9abb79ec5/defusedxml-0.8.0rc2-py2.py3-none-any.whl", hash = "sha256:1c812964311154c3bf4aaf3bc1443b31ee13530b7f255eaaa062c0553c76103d", size = 25756, upload-time = "2023-09-29T08:01:25.515Z" },
-]
-
 [[package]]
 name = "deprecated"
 version = "1.2.18"
@@ -1112,16 +1107,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/8f/d7/9322c609343d929e75e7e5e6255e614fcc67572cfd083959cdef3b7aad79/docutils-0.21.2-py3-none-any.whl", hash = "sha256:dafca5b9e384f0e419294eb4d2ff9fa826435bf15f15b7bd45723e8ad76811b2", size = 587408, upload-time = "2024-04-23T18:57:14.835Z" },
 ]
 
-[[package]]
-name = "drain3"
-version = "0.9.11"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "cachetools" },
-    { name = "jsonpickle" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/dc/83/4da2d3a11b5e0edf1a4f4c0c2dd42126d2eb1f31c733967edd3dfac1af94/drain3-0.9.11.tar.gz", hash = "sha256:9ab4b1407fad74f56554ae371ef019c3c7985861631f4bab46a0e92585125f75", size = 27960, upload-time = "2022-07-17T06:40:11.433Z" }
-
 [[package]]
 name = "dulwich"
 version = "0.25.2"
@@ -1709,30 +1694,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ad/e7/d6914822c88aa2974dbbd10903d801a28a19ce9cd8bad7e694cbbcf61528/grpcio-1.78.0-cp312-cp312-win_amd64.whl", hash = "sha256:9566fe4ababbb2610c39190791e5b829869351d14369603702e890ef3ad2d06e", size = 4797657, upload-time = "2026-02-06T09:55:49.86Z" },
 ]
 
-[[package]]
-name = "grpcio-tools"
-version = "1.78.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "grpcio" },
-    { name = "protobuf" },
-    { name = "setuptools", version = "79.0.1", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
-    { name = "setuptools", version = "80.10.2", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'darwin' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/8b/d1/cbefe328653f746fd319c4377836a25ba64226e41c6a1d7d5cdbc87a459f/grpcio_tools-1.78.0.tar.gz", hash = "sha256:4b0dd86560274316e155d925158276f8564508193088bc43e20d3f5dff956b2b", size = 5393026, upload-time = "2026-02-06T09:59:59.53Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/0c/ae/5b1fa5dd8d560a6925aa52de0de8731d319f121c276e35b9b2af7cc220a2/grpcio_tools-1.78.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:9eb122da57d4cad7d339fc75483116f0113af99e8d2c67f3ef9cae7501d806e4", size = 2546823, upload-time = "2026-02-06T09:58:17.944Z" },
-    { url = "https://files.pythonhosted.org/packages/a7/ed/d33ccf7fa701512efea7e7e23333b748848a123e9d3bbafde4e126784546/grpcio_tools-1.78.0-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:d0c501b8249940b886420e6935045c44cb818fa6f265f4c2b97d5cff9cb5e796", size = 5706776, upload-time = "2026-02-06T09:58:20.944Z" },
-    { url = "https://files.pythonhosted.org/packages/c6/69/4285583f40b37af28277fc6b867d636e3b10e1b6a7ebd29391a856e1279b/grpcio_tools-1.78.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:77e5aa2d2a7268d55b1b113f958264681ef1994c970f69d48db7d4683d040f57", size = 2593972, upload-time = "2026-02-06T09:58:23.29Z" },
-    { url = "https://files.pythonhosted.org/packages/d7/eb/ecc1885bd6b3147f0a1b7dff5565cab72f01c8f8aa458f682a1c77a9fb08/grpcio_tools-1.78.0-cp312-cp312-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:8e3c0b0e6ba5275322ba29a97bf890565a55f129f99a21b121145e9e93a22525", size = 2905531, upload-time = "2026-02-06T09:58:25.406Z" },
-    { url = "https://files.pythonhosted.org/packages/ae/a9/511d0040ced66960ca10ba0f082d6b2d2ee6dd61837b1709636fdd8e23b4/grpcio_tools-1.78.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:975d4cb48694e20ebd78e1643e5f1cd94cdb6a3d38e677a8e84ae43665aa4790", size = 2656909, upload-time = "2026-02-06T09:58:28.022Z" },
-    { url = "https://files.pythonhosted.org/packages/06/a3/3d2c707e7dee8df842c96fbb24feb2747e506e39f4a81b661def7fed107c/grpcio_tools-1.78.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:553ff18c5d52807dedecf25045ae70bad7a3dbba0b27a9a3cdd9bcf0a1b7baec", size = 3109778, upload-time = "2026-02-06T09:58:30.091Z" },
-    { url = "https://files.pythonhosted.org/packages/1f/4b/646811ba241bf05da1f0dc6f25764f1c837f78f75b4485a4210c84b79eae/grpcio_tools-1.78.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:8c7f5e4af5a84d2e96c862b1a65e958a538237e268d5f8203a3a784340975b51", size = 3658763, upload-time = "2026-02-06T09:58:32.875Z" },
-    { url = "https://files.pythonhosted.org/packages/45/de/0a5ef3b3e79d1011375f5580dfee3a9c1ccb96c5f5d1c74c8cee777a2483/grpcio_tools-1.78.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:96183e2b44afc3f9a761e9d0f985c3b44e03e8bb98e626241a6cbfb3b6f7e88f", size = 3325116, upload-time = "2026-02-06T09:58:34.894Z" },
-    { url = "https://files.pythonhosted.org/packages/95/d2/6391b241ad571bc3e71d63f957c0b1860f0c47932d03c7f300028880f9b8/grpcio_tools-1.78.0-cp312-cp312-win32.whl", hash = "sha256:2250e8424c565a88573f7dc10659a0b92802e68c2a1d57e41872c9b88ccea7a6", size = 993493, upload-time = "2026-02-06T09:58:37.242Z" },
-    { url = "https://files.pythonhosted.org/packages/7c/8f/7d0d3a39ecad76ccc136be28274daa660569b244fa7d7d0bbb24d68e5ece/grpcio_tools-1.78.0-cp312-cp312-win_amd64.whl", hash = "sha256:217d1fa29de14d9c567d616ead7cb0fef33cde36010edff5a9390b00d52e5094", size = 1158423, upload-time = "2026-02-06T09:58:40.072Z" },
-]
-
 [[package]]
 name = "gunicorn"
 version = "23.0.0"
@@ -2081,36 +2042,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/1e/e8/685f47e0d754320684db4425a0967f7d3fa70126bffd76110b7009a0090f/joblib-1.5.2-py3-none-any.whl", hash = "sha256:4e1f0bdbb987e6d843c70cf43714cb276623def372df3c22fe5266b2670bc241", size = 308396, upload-time = "2025-08-27T12:15:45.188Z" },
 ]
 
-[[package]]
-name = "jsonpatch"
-version = "1.33"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "jsonpointer" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/42/78/18813351fe5d63acad16aec57f94ec2b70a09e53ca98145589e185423873/jsonpatch-1.33.tar.gz", hash = "sha256:9fcd4009c41e6d12348b4a0ff2563ba56a2923a7dfee731d004e212e1ee5030c", size = 21699, upload-time = "2023-06-26T12:07:29.144Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/73/07/02e16ed01e04a374e644b575638ec7987ae846d25ad97bcc9945a3ee4b0e/jsonpatch-1.33-py2.py3-none-any.whl", hash = "sha256:0ae28c0cd062bbd8b8ecc26d7d164fbbea9652a1a3693f3b956c1eae5145dade", size = 12898, upload-time = "2023-06-16T21:01:28.466Z" },
-]
-
-[[package]]
-name = "jsonpickle"
-version = "1.5.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/af/62/31ef0b58050a3731b079af69932104a9443bff07fe2b9564c161e3ec4348/jsonpickle-1.5.1.tar.gz", hash = "sha256:060f97096559d1b86aa16cac2f4ea5f7b6da0c15d8a4de150b78013a886f9a51", size = 109560, upload-time = "2021-01-31T05:57:15.037Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/77/a7/c2f527ddce3155ae9e008385963c2325cbfd52969f8b38efa2723e2af4af/jsonpickle-1.5.1-py2.py3-none-any.whl", hash = "sha256:8eb8323f0e12cb40687f0445e2115d8165901e20ac670add55bb53a95c68c0e5", size = 37124, upload-time = "2021-01-31T05:57:12.256Z" },
-]
-
-[[package]]
-name = "jsonpointer"
-version = "3.1.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/18/c7/af399a2e7a67fd18d63c40c5e62d3af4e67b836a2107468b6a5ea24c4304/jsonpointer-3.1.1.tar.gz", hash = "sha256:0b801c7db33a904024f6004d526dcc53bbb8a4a0f4e32bfd10beadf60adf1900", size = 9068, upload-time = "2026-03-23T22:32:32.458Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/9e/6a/a83720e953b1682d2d109d3c2dbb0bc9bf28cc1cbc205be4ef4be5da709d/jsonpointer-3.1.1-py3-none-any.whl", hash = "sha256:8ff8b95779d071ba472cf5bc913028df06031797532f08a7d5b602d8b2a488ca", size = 7659, upload-time = "2026-03-23T22:32:31.568Z" },
-]
-
 [[package]]
 name = "jsonschema"
 version = "4.25.1"
@@ -2159,103 +2090,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ff/fa/538442202d639add2f52a814bdfc58207ee6fbb6d1ecd1a6e867f48ec1af/kiwisolver-1.4.10rc0-cp312-cp312-win_arm64.whl", hash = "sha256:44cd6dfea8a6c2becac4f3d60ebdcfe4fed858bbf7fe9cd38ffea7b58df66435", size = 65077, upload-time = "2025-08-10T20:21:01.882Z" },
 ]
 
-[[package]]
-name = "langchain"
-version = "0.3.29"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "langchain-core" },
-    { name = "langchain-text-splitters" },
-    { name = "langsmith" },
-    { name = "pydantic" },
-    { name = "pyyaml" },
-    { name = "requests" },
-    { name = "sqlalchemy" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/89/b8/b07646e1d0caa372dcd2a9bb9c0eaa94ba5cc3517107b97f4f4f348b04ec/langchain-0.3.29.tar.gz", hash = "sha256:4c7da1879c4138de4721ce4e9dfbda97fa0a1a4afb24f61bcfe6643a031e99fe", size = 10244120, upload-time = "2026-05-05T21:01:45.027Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/6c/c8/36c31c6b50715e04b5a3bfd0491fb49ce9bb2c098724c042c9a9450f1c6c/langchain-0.3.29-py3-none-any.whl", hash = "sha256:4af0c6056f2ace34813b77412cceb885b419e2c43d225742ae01f4a32171ccc9", size = 1026012, upload-time = "2026-05-05T21:01:42.187Z" },
-]
-
-[[package]]
-name = "langchain-core"
-version = "0.3.85"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "jsonpatch" },
-    { name = "langsmith" },
-    { name = "packaging" },
-    { name = "pydantic" },
-    { name = "pyyaml" },
-    { name = "tenacity" },
-    { name = "typing-extensions" },
-    { name = "uuid-utils" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/24/48/fd8dee0e2bf19f94d9c04d62131b9dbf51f7d08e645dd00daa5d6c9eb60e/langchain_core-0.3.85.tar.gz", hash = "sha256:9321142eb8f754045c02b914bc9a18b07a589e76423c3ac0e69c35e130826f0c", size = 601850, upload-time = "2026-05-05T20:43:20.865Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/b1/71/cadea475bda84fd4cc6fe46b347fd643be8b323fc8c413f4ccbbe2717db6/langchain_core-0.3.85-py3-none-any.whl", hash = "sha256:3ff7d8da9645cc8ff221d3db7fcccb9794ea4b84834b3714071459c254ed9061", size = 460237, upload-time = "2026-05-05T20:43:19.294Z" },
-]
-
-[[package]]
-name = "langchain-nvidia-ai-endpoints"
-version = "0.3.19"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "aiohttp" },
-    { name = "filetype" },
-    { name = "langchain-core" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/4d/6f/0fcad6732f124c0d96fd7e8487a0b51db9f167964cfe041e4669a2148dbe/langchain_nvidia_ai_endpoints-0.3.19.tar.gz", hash = "sha256:1c420c88f7c78b2b2c2fdcef8e46104c2dd19c81e036bdd03a4838a6340950fe", size = 42884, upload-time = "2025-10-31T00:17:19.352Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/de/46/d7529004de384b2abc9e5b76cf4a84a23f3028ec6381bd5f7c00ac39bfab/langchain_nvidia_ai_endpoints-0.3.19-py3-none-any.whl", hash = "sha256:40161a71646fcbe457ac5f2222c5eadcbe31a7d79d618f5a0857c37fffa3a6d5", size = 46229, upload-time = "2025-10-31T00:17:18.306Z" },
-]
-
-[[package]]
-name = "langchain-openai"
-version = "0.3.35"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "langchain-core" },
-    { name = "openai" },
-    { name = "tiktoken" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/fb/96/06d0d25a37e05a0ff2d918f0a4b0bf0732aed6a43b472b0b68426ce04ef8/langchain_openai-0.3.35.tar.gz", hash = "sha256:fa985fd041c3809da256a040c98e8a43e91c6d165b96dcfeb770d8bd457bf76f", size = 786635, upload-time = "2025-10-06T15:09:28.463Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/d8/d5/c90c5478215c20ee71d8feaf676f7ffd78d0568f8c98bd83f81ce7562ed7/langchain_openai-0.3.35-py3-none-any.whl", hash = "sha256:76d5707e6e81fd461d33964ad618bd326cb661a1975cef7c1cb0703576bdada5", size = 75952, upload-time = "2025-10-06T15:09:27.137Z" },
-]
-
-[[package]]
-name = "langchain-text-splitters"
-version = "0.3.11"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "langchain-core" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/11/43/dcda8fd25f0b19cb2835f2f6bb67f26ad58634f04ac2d8eae00526b0fa55/langchain_text_splitters-0.3.11.tar.gz", hash = "sha256:7a50a04ada9a133bbabb80731df7f6ddac51bc9f1b9cab7fa09304d71d38a6cc", size = 46458, upload-time = "2025-08-31T23:02:58.316Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/58/0d/41a51b40d24ff0384ec4f7ab8dd3dcea8353c05c973836b5e289f1465d4f/langchain_text_splitters-0.3.11-py3-none-any.whl", hash = "sha256:cf079131166a487f1372c8ab5d0bfaa6c0a4291733d9c43a34a16ac9bcd6a393", size = 33845, upload-time = "2025-08-31T23:02:57.195Z" },
-]
-
-[[package]]
-name = "langsmith"
-version = "0.8.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "httpx" },
-    { name = "orjson", marker = "platform_python_implementation != 'PyPy' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
-    { name = "packaging" },
-    { name = "pydantic" },
-    { name = "requests" },
-    { name = "requests-toolbelt" },
-    { name = "uuid-utils" },
-    { name = "xxhash" },
-    { name = "zstandard" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/22/f1/717e01ab46dd397f4ed19bb230c045b5e80ec2c0eedb42941b2b68d07032/langsmith-0.8.1.tar.gz", hash = "sha256:63171ca4fccd6a3209539a7fef4d0e7edc6437d142f6740a6a383bee911bd17e", size = 4457870, upload-time = "2026-05-05T20:08:58.716Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/5d/2c/279ad7b6acff0704fa66ee52e4f66669fe948df6502bd5982b53d3612c06/langsmith-0.8.1-py3-none-any.whl", hash = "sha256:8809f43d44d53ac3f21127f61fff7f8bbc23e64f164c29d2df8c475ec41be6c3", size = 397537, upload-time = "2026-05-05T20:08:56.808Z" },
-]
-
 [[package]]
 name = "lark"
 version = "1.2.2"
@@ -2361,25 +2195,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a0/ef/11292bb0b85cf4c93447cab5a29f64576ed14d3ab4280e35ddd23486594a/lm_format_enforcer-0.11.3-py3-none-any.whl", hash = "sha256:cf586350875def1ae7a8fba84fcbbfc8371424b6c9d05c1fcba70aa233fbf06f", size = 45418, upload-time = "2025-08-24T19:37:46.325Z" },
 ]
 
-[[package]]
-name = "logsage"
-version = "0.1.7"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "drain3" },
-    { name = "langchain" },
-    { name = "langchain-core" },
-    { name = "langchain-nvidia-ai-endpoints" },
-    { name = "nh3" },
-    { name = "numpy" },
-    { name = "pandas" },
-    { name = "pydantic-settings" },
-    { name = "requests" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/00/29/41ca46b94399d55569b1a19d909115cfef47456b88bb302960d58e3fd1f3/logsage-0.1.7-py3-none-any.whl", hash = "sha256:690e9f6dc56bf369b90aad91d7463e4c1689feb589148481955437ec8d33088a", size = 75267, upload-time = "2026-04-13T07:59:47.001Z" },
-]
-
 [[package]]
 name = "loguru"
 version = "0.7.3"
@@ -2482,14 +2297,14 @@ dependencies = [
     { name = "jsonschema" },
     { name = "pydantic" },
     { name = "pydantic-settings" },
-    { name = "pyjwt", extra = ["crypto"] },
+    { name = "pyjwt", extra = ["crypto"], marker = "extra == 'extra-18-nemo-export-deploy-vllm'" },
     { name = "python-multipart" },
-    { name = "pywin32", marker = "sys_platform == 'win32' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "pywin32", marker = "sys_platform == 'win32'" },
     { name = "sse-starlette" },
     { name = "starlette" },
     { name = "typing-extensions" },
     { name = "typing-inspection" },
-    { name = "uvicorn", marker = "sys_platform != 'emscripten' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
+    { name = "uvicorn", marker = "sys_platform != 'emscripten'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/fc/6d/62e76bbb8144d6ed86e202b5edd8a4cb631e7c8130f3f4893c3f90262b10/mcp-1.26.0.tar.gz", hash = "sha256:db6e2ef491eecc1a0d93711a76f28dec2e05999f93afd48795da1c1137142c66", size = 608005, upload-time = "2026-01-24T19:40:32.468Z" }
 wheels = [
@@ -3096,30 +2911,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/eb/8d/776adee7bbf76365fdd7f2552710282c79a4ead5d2a46408c9043a2b70ba/networkx-3.5-py3-none-any.whl", hash = "sha256:0030d386a9a06dee3565298b4a734b68589749a544acbb6c412dc9e2489ec6ec", size = 2034406, upload-time = "2025-05-29T11:35:04.961Z" },
 ]
 
-[[package]]
-name = "nh3"
-version = "0.3.5"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/9c/5f/1d19bdc7d27238e37f3672cdc02cb77c56a4a86d140cd4f4f23c90df6e16/nh3-0.3.5.tar.gz", hash = "sha256:45855e14ff056064fec77133bfcf7cd691838168e5e17bbef075394954dc9dc8", size = 20743, upload-time = "2026-04-25T10:44:16.066Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/85/30/d162e99746a2fb1d98bb0ef23af3e201b156cf09f7de867c7390c8fe1c06/nh3-0.3.5-cp38-abi3-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:3bb854485c9b33e5bb143ff3e49e577073bc6bc320f0ff8fc316dd89c0d3c101", size = 1442393, upload-time = "2026-04-25T10:43:53.556Z" },
-    { url = "https://files.pythonhosted.org/packages/25/8c/072120d506978ab053e1732d0efa7c86cb478fee0ee098fda0ac0d31cb34/nh3-0.3.5-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:50d401ab2d8e86d59e2126e3ab2a2f45840c405842b626d9a51624b3a33b6878", size = 837722, upload-time = "2026-04-25T10:43:55.073Z" },
-    { url = "https://files.pythonhosted.org/packages/52/86/d4e06e28c5ad1c4b065f89737d02631bd49f1660b6ebcf17a87ffcd201da/nh3-0.3.5-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:acfd354e61accbe4c74f8017c6e397a776916dfe47c48643cf7fd84ade826f93", size = 822872, upload-time = "2026-04-25T10:43:56.581Z" },
-    { url = "https://files.pythonhosted.org/packages/0a/62/50659255213f241ec5797ae7427464c969397373e83b3659372b341ae869/nh3-0.3.5-cp38-abi3-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:52d877980d7ca01dc3baf3936bf844828bc6f332962227a684ed79c18cce14c3", size = 1100031, upload-time = "2026-04-25T10:43:58.098Z" },
-    { url = "https://files.pythonhosted.org/packages/00/7a/a12ae77593b2fcf3be25df7bc1c01967d0de448bdb4b6c7ec80fe4f5a74f/nh3-0.3.5-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:207c01801d3e9bb8ec08f08689346bdd30ce15b8bf60013a925d08b5388962a4", size = 1057669, upload-time = "2026-04-25T10:43:59.328Z" },
-    { url = "https://files.pythonhosted.org/packages/2d/71/5647dc04c0233192a3956fc91708822b21403a06508cacf78083c68e7bf0/nh3-0.3.5-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ea232933394d1d58bf7c4bb348dc4660eae6604e1ae81cd2ba6d9ed80d390f3b", size = 914795, upload-time = "2026-04-25T10:44:00.52Z" },
-    { url = "https://files.pythonhosted.org/packages/1b/0e/bf298920729f216adcb002acf7ea01b90842603d2e4e2ce9b900d9ee8fab/nh3-0.3.5-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fe3a787dc76b50de6bee54ef242f26c41dfe47654428e3e94f0fae5bb6dd2cc1", size = 806976, upload-time = "2026-04-25T10:44:01.743Z" },
-    { url = "https://files.pythonhosted.org/packages/85/01/26761e1dc2b848e65a62c19e5d39ad446283287cd4afddc89f364ab86bc9/nh3-0.3.5-cp38-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:488928988caad25ba14b1eb5bc74e25e21f3b5e40341d956f3ce4a8bc19460dc", size = 834904, upload-time = "2026-04-25T10:44:03.454Z" },
-    { url = "https://files.pythonhosted.org/packages/33/53/0766113e679540ac1edc1b82b1295aecd321eeb75d6fead70109a838b6ee/nh3-0.3.5-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:2c069570b06aa848457713ad7af4a9905691291548c4466a9ad78ee95808382b", size = 857159, upload-time = "2026-04-25T10:44:05.003Z" },
-    { url = "https://files.pythonhosted.org/packages/58/36/734d353dfaf292fed574b8b3092f0ef79dc6404f3879f7faaa61a4701fad/nh3-0.3.5-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:eeedc90ed8c42c327e8e10e621ccfa314fc6cce35d5929f4297ff1cdb89667c4", size = 1018600, upload-time = "2026-04-25T10:44:06.18Z" },
-    { url = "https://files.pythonhosted.org/packages/6b/aa/d9c59c1b49669fcb7bababa55df82385f029ad5c2651f583c3a1141cfdd1/nh3-0.3.5-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:de8e8621853b6470fe928c684ee0d3f39ea8086cebafe4c416486488dea7b68d", size = 1103530, upload-time = "2026-04-25T10:44:07.68Z" },
-    { url = "https://files.pythonhosted.org/packages/90/b0/cdd210bfb8d9d43fb02fc3c868336b9955934d8e15e66eb1d15a147b8af0/nh3-0.3.5-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:6ea58cc44d274c643b83547ca9654a0b1a817609b160601356f76a2b744c49ad", size = 1061754, upload-time = "2026-04-25T10:44:09.362Z" },
-    { url = "https://files.pythonhosted.org/packages/ce/cb/7a39e72e668c8445bdd95e494b3e21cfdddc68329be8ea3522c8befb46c4/nh3-0.3.5-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e49c9b564e6bcb03ecd2f057213df9a0de15a95812ac9db9600b590db23d3ae9", size = 1040938, upload-time = "2026-04-25T10:44:10.775Z" },
-    { url = "https://files.pythonhosted.org/packages/af/4c/fc2f9ed208a3801a319f59b5fea03cdc20cf3bd8af14be930d3a8de01224/nh3-0.3.5-cp38-abi3-win32.whl", hash = "sha256:559e4c73b689e9a7aa97ac9760b1bc488038d7c1a575aa4ab5a0e19ee9630c0f", size = 611445, upload-time = "2026-04-25T10:44:12.317Z" },
-    { url = "https://files.pythonhosted.org/packages/db/1a/e4c9b5e2ae13e6092c9ec16d8ca30646cb01fcdea245f36c5b08fd21fbd5/nh3-0.3.5-cp38-abi3-win_amd64.whl", hash = "sha256:45e6a65dc88a300a2e3502cb9c8e6d1d6b831d6fba7470643333609c6aab1f30", size = 626502, upload-time = "2026-04-25T10:44:13.682Z" },
-    { url = "https://files.pythonhosted.org/packages/80/7c/19cd0671d1ba2762fb388fc149697d20d0568ccfeef833b11280a619e526/nh3-0.3.5-cp38-abi3-win_arm64.whl", hash = "sha256:8f85285700a18e9f3fc5bff41fe573fa84f81542ef13b48a89f9fecca0474d3b", size = 611069, upload-time = "2026-04-25T10:44:14.934Z" },
-]
-
 [[package]]
 name = "ninja"
 version = "1.13.0"
@@ -3469,23 +3260,8 @@ wheels = [
 
 [[package]]
 name = "nvidia-resiliency-ext"
-version = "0.6.0.dev69+b2bb3d7"
+version = "0.6.0"
 source = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git?rev=b2bb3d728a18795807d9f76c535e005a609a1b01#b2bb3d728a18795807d9f76c535e005a609a1b01" }
-dependencies = [
-    { name = "defusedxml" },
-    { name = "grpcio" },
-    { name = "grpcio-tools" },
-    { name = "langchain-openai" },
-    { name = "logsage" },
-    { name = "mcp" },
-    { name = "nvidia-ml-py" },
-    { name = "packaging" },
-    { name = "protobuf" },
-    { name = "psutil" },
-    { name = "pyyaml" },
-    { name = "setproctitle" },
-    { name = "torch", marker = "sys_platform == 'never' or (extra == 'extra-18-nemo-export-deploy-trt-onnx' and extra == 'extra-18-nemo-export-deploy-vllm')" },
-]
 
 [[package]]
 name = "nvidia-sphinx-theme"
@@ -5385,15 +5161,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/40/44/4a5f08c96eb108af5cb50b41f76142f0afa346dfa99d5296fe7202a11854/tabulate-0.9.0-py3-none-any.whl", hash = "sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f", size = 35252, upload-time = "2022-10-06T17:21:44.262Z" },
 ]
 
-[[package]]
-name = "tenacity"
-version = "9.1.4"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/47/c6/ee486fd809e357697ee8a44d3d69222b344920433d3b6666ccd9b374630c/tenacity-9.1.4.tar.gz", hash = "sha256:adb31d4c263f2bd041081ab33b498309a57c77f9acf2db65aadf0898179cf93a", size = 49413, upload-time = "2026-02-07T10:45:33.841Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl", hash = "sha256:6095a360c919085f28c6527de529e76a06ad89b23659fa881ae0649b867a9d55", size = 28926, upload-time = "2026-02-07T10:45:32.24Z" },
-]
-
 [[package]]
 name = "tensorboard"
 version = "2.20.0"
@@ -5865,28 +5632,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a7/c2/fe1e52489ae3122415c51f387e221dd0773709bad6c6cdaa599e8a2c5185/urllib3-2.5.0-py3-none-any.whl", hash = "sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc", size = 129795, upload-time = "2025-06-18T14:07:40.39Z" },
 ]
 
-[[package]]
-name = "uuid-utils"
-version = "0.15.0a3"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/84/49/1010867ba865e0e35bb70c8fa04698975eecc132bc336e557f72fa01ab19/uuid_utils-0.15.0a3.tar.gz", hash = "sha256:5122bae6ac698b9e53be18d3d54742350a796094d68ebadb4999f55de4901d5f", size = 25663, upload-time = "2026-04-16T17:49:05.704Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/80/36/9a699d701a931a6026ced2a9005735e4755299bdc39e4368f75bfd31f696/uuid_utils-0.15.0a3-cp312-cp312-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:5e623d0f221da43a3ed41e5e7b12405d572be3ef346159f02651027161bbaa07", size = 559521, upload-time = "2026-04-16T17:50:17.966Z" },
-    { url = "https://files.pythonhosted.org/packages/86/34/c8c9545dceba811ceaa2bddfd5e2dde05110cd502ec47cc647caed8ef383/uuid_utils-0.15.0a3-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:80bf22c8901b50f9c1f43b932657a645a449cefbb3e0c80b9c53dbe7afb08e1d", size = 288287, upload-time = "2026-04-16T17:50:35.322Z" },
-    { url = "https://files.pythonhosted.org/packages/3c/75/57ee912882c9406ea775f9a4e4a7203e5f2874ad1049e0cb455d45c69afa/uuid_utils-0.15.0a3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fc22cafdffd573bd76f8fc0ff1140109c9ba26fe6c99fba2b0f497c2c385f40f", size = 324633, upload-time = "2026-04-16T17:50:29.937Z" },
-    { url = "https://files.pythonhosted.org/packages/07/a9/e03bdedf8855d80cb152b9f2bafed6762a0c7a0b23d6ed8a1ba37b60ddfa/uuid_utils-0.15.0a3-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:d395ce1db61ef8040ea60a1fe03b6af6f91b09a0ec811ae18db75202278267cc", size = 331071, upload-time = "2026-04-16T17:48:29.65Z" },
-    { url = "https://files.pythonhosted.org/packages/46/5b/6294519d94bfb129dd7a2fabadfa9f332bc556dc9bcc0d65a7f88212c356/uuid_utils-0.15.0a3-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b75ca2b2ce7fb4e9107c4e04d388abea0bfa3454c6735e4bb69dae31bae7ef82", size = 443936, upload-time = "2026-04-16T17:49:23.365Z" },
-    { url = "https://files.pythonhosted.org/packages/9c/68/b903b7198ea6e402a95e8732795cf7f2fd59976d1129da94376f865be83c/uuid_utils-0.15.0a3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7b081c07f31b82db60ac9f2205f2dcdabfc596cc07efc82b799fa3f7212e492e", size = 324611, upload-time = "2026-04-16T17:49:28.467Z" },
-    { url = "https://files.pythonhosted.org/packages/61/a0/00cd1db7ff77f24ac36d34397f02e2fbc5f7005323be8a319e6f459467a2/uuid_utils-0.15.0a3-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c7892ffca114a3a9074713f7b449a5c3fde69c849d941d5bafd7a99e2716229f", size = 348374, upload-time = "2026-04-16T17:50:50.631Z" },
-    { url = "https://files.pythonhosted.org/packages/1c/4c/2bff6329e8307cf57eacfdf384294e51b8926485627c21956eaaa26fba3b/uuid_utils-0.15.0a3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d64b2fb785a1d84f64c5bd45cf391316f73f226ec3db8379215862b8f4a11bb3", size = 501052, upload-time = "2026-04-16T17:48:38.631Z" },
-    { url = "https://files.pythonhosted.org/packages/75/7a/0be7b48c187aeb16bd31928f5b32a2050824759053c2bb4880b29ea840b8/uuid_utils-0.15.0a3-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:cc47103df3816ab6619ea022f8c07f49be0cc8901588a99d87b294b78894694c", size = 606361, upload-time = "2026-04-16T17:50:00.019Z" },
-    { url = "https://files.pythonhosted.org/packages/07/c4/7fc8b92eb1d6f0f5ac784a7bdca3818e708580f8add0f7bf72b6918ca6d0/uuid_utils-0.15.0a3-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:d9b7a8e278c827739cbe22e6debce6bf7d4cb3887f33987294ca45b503994d29", size = 564846, upload-time = "2026-04-16T17:49:24.893Z" },
-    { url = "https://files.pythonhosted.org/packages/90/86/1e2baaa34067301af49bd6b46133c69b99134babba3ac1db56129ea14648/uuid_utils-0.15.0a3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:923f1bd9f8b75e9f65d8cb2b442d85ee57cff55654eddd92e230c59e70fdc9b6", size = 529734, upload-time = "2026-04-16T17:49:21.913Z" },
-    { url = "https://files.pythonhosted.org/packages/18/b0/d8039a896b2d0365842dc614bb12067e81d312fc98dc95d816ed5bd2e4f4/uuid_utils-0.15.0a3-cp312-cp312-win32.whl", hash = "sha256:27d1604f422ff9c9627821963dd01c86dc40397da0bd43d154d7ffe6806cc050", size = 168170, upload-time = "2026-04-16T17:50:34.029Z" },
-    { url = "https://files.pythonhosted.org/packages/fc/b5/19a451abacb2ba47faec3cc6409845427d496a5afd0f69327db21d8c30b1/uuid_utils-0.15.0a3-cp312-cp312-win_amd64.whl", hash = "sha256:375ff8a31513e690b95ca7902f912b00315b0f6d0022b708a9e315529b4df75d", size = 173937, upload-time = "2026-04-16T17:49:08.148Z" },
-    { url = "https://files.pythonhosted.org/packages/69/91/53728a4a94823b39e9c81b04d1d41547de252e3a3a66df28a6ca165aaf64/uuid_utils-0.15.0a3-cp312-cp312-win_arm64.whl", hash = "sha256:2b4697e395907895c6e945865103ff24e584f866cbd5802d56bcdf2bab5c4312", size = 172214, upload-time = "2026-04-16T17:50:15.338Z" },
-]
-
 [[package]]
 name = "uvicorn"
 version = "0.37.0"

From f63b3ced38fa00d0313ca2b50213573eccaa10fb Mon Sep 17 00:00:00 2001
From: Onur Yilmaz <oyilmaz@nvidia.com>
Date: Wed, 1 Jul 2026 09:26:38 -0400
Subject: [PATCH 4/5] Pass trust param for tokenizer

Signed-off-by: Onur Yilmaz <oyilmaz@nvidia.com>
---
 nemo_deploy/llm/inference/inference_base.py     | 5 ++++-
 nemo_deploy/llm/megatronllm_deployable.py       | 2 ++
 pyproject.toml                                  | 1 +
 scripts/deploy/llm/mlm/deploy_triton.py         | 8 ++++++++
 scripts/deploy/nlp/deploy_inframework_triton.py | 8 ++++++++
 scripts/deploy/nlp/deploy_ray_inframework.py    | 8 ++++++++
 6 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/nemo_deploy/llm/inference/inference_base.py b/nemo_deploy/llm/inference/inference_base.py
index b2f16d4fd..30aeab7ce 100644
--- a/nemo_deploy/llm/inference/inference_base.py
+++ b/nemo_deploy/llm/inference/inference_base.py
@@ -194,6 +194,7 @@ def setup_megatron_model_and_tokenizer_for_inference(
     sequence_parallel: Optional[bool] = None,
     micro_batch_size: Optional[int] = None,
     model_type: str = "gpt",
+    trust_remote_code: bool = False,
 ) -> Tuple[List[MegatronModule], MegatronTokenizer]:
     """Initialize a Megatron model and tokenizer for inference from a Megatron-LM/MBridge checkpoint.
 
@@ -259,7 +260,7 @@ def setup_megatron_model_and_tokenizer_for_inference(
         megatron_args=mlm_args,
         use_cpu_init=False,
     )
-    tokenizer = load_tokenizer(checkpoint_path)
+    tokenizer = load_tokenizer(checkpoint_path, trust_remote_code=trust_remote_code)
     return model, tokenizer, mlm_args
 
 
@@ -436,6 +437,7 @@ def create_mcore_engine(
     micro_batch_size: Optional[int] = None,
     buffer_size_gb: float = 10.0,
     legacy_model_format: bool = False,
+    trust_remote_code: bool = False,
     **model_config_kwargs,
 ) -> Tuple[MCoreEngineWithCleanup, Union[MCoreTokenizerWrappper, MegatronTokenizer]]:
     """Set up the model, tokenizer and MegatronLLM engine for inference.
@@ -497,6 +499,7 @@ def create_mcore_engine(
             sequence_parallel=sequence_parallel,
             micro_batch_size=micro_batch_size,
             model_type=model_type,
+            trust_remote_code=trust_remote_code,
         )
         model = modelList[0]
     else:
diff --git a/nemo_deploy/llm/megatronllm_deployable.py b/nemo_deploy/llm/megatronllm_deployable.py
index 10910ac90..5791ab335 100755
--- a/nemo_deploy/llm/megatronllm_deployable.py
+++ b/nemo_deploy/llm/megatronllm_deployable.py
@@ -105,6 +105,7 @@ def __init__(
         micro_batch_size: Optional[int] = None,
         buffer_size_gb: float = 10.0,
         legacy_model_format: bool = False,
+        trust_remote_code: bool = False,
         **model_config_kwargs,
     ):
         if not HAVE_TRITON:
@@ -136,6 +137,7 @@ def __init__(
             micro_batch_size=micro_batch_size,
             buffer_size_gb=buffer_size_gb,
             legacy_model_format=legacy_model_format,
+            trust_remote_code=trust_remote_code,
             **model_config_kwargs,
         )
         self.enable_cuda_graphs = enable_cuda_graphs
diff --git a/pyproject.toml b/pyproject.toml
index c0d5666d4..842ff99d1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -123,6 +123,7 @@ no-build-isolation-package = [
     "mamba-ssm",
     "causal-conv1d",
     "nv-grouped-gemm",
+    "fast-hadamard-transform",
 ]
 
 # Always apply the build group since dependencies like TE/mcore/nemo-run require build dependencies
diff --git a/scripts/deploy/llm/mlm/deploy_triton.py b/scripts/deploy/llm/mlm/deploy_triton.py
index 3fd0cdcda..51a347c1a 100755
--- a/scripts/deploy/llm/mlm/deploy_triton.py
+++ b/scripts/deploy/llm/mlm/deploy_triton.py
@@ -216,6 +216,13 @@ def get_args(argv):
         default=None,
         help="Micro batch size for model execution",
     )
+    parser.add_argument(
+        "-trc",
+        "--trust_remote_code",
+        default=False,
+        action="store_true",
+        help="Pass trust_remote_code=True to load_tokenizer() for checkpoints whose tokenizer requires it.",
+    )
     args = parser.parse_args(argv)
     return args
 
@@ -263,6 +270,7 @@ def nemo_deploy(argv):
         legacy_ckpt=args.legacy_ckpt,
         model_type=args.model_type,
         micro_batch_size=args.micro_batch_size,
+        trust_remote_code=args.trust_remote_code,
         **model_config_kwargs,
     )
 
diff --git a/scripts/deploy/nlp/deploy_inframework_triton.py b/scripts/deploy/nlp/deploy_inframework_triton.py
index 3db1a46da..1f86d0143 100755
--- a/scripts/deploy/nlp/deploy_inframework_triton.py
+++ b/scripts/deploy/nlp/deploy_inframework_triton.py
@@ -234,6 +234,13 @@ def get_args(argv):
         action="store_true",
         help="Use the legacy StaticInferenceEngine path in MCoreEngine",
     )
+    parser.add_argument(
+        "-trc",
+        "--trust_remote_code",
+        default=False,
+        action="store_true",
+        help="Pass trust_remote_code=True to load_tokenizer() for checkpoints whose tokenizer requires it.",
+    )
     args = parser.parse_args(argv)
     return args
 
@@ -284,6 +291,7 @@ def nemo_deploy(argv):
         model_type=args.model_type,
         micro_batch_size=args.micro_batch_size,
         legacy_model_format=args.legacy_model_format,
+        trust_remote_code=args.trust_remote_code,
         **model_config_kwargs,
     )
 
diff --git a/scripts/deploy/nlp/deploy_ray_inframework.py b/scripts/deploy/nlp/deploy_ray_inframework.py
index 23d701f26..cbabbbc26 100644
--- a/scripts/deploy/nlp/deploy_ray_inframework.py
+++ b/scripts/deploy/nlp/deploy_ray_inframework.py
@@ -225,6 +225,13 @@ def parse_args():
         default={},
         help="Runtime environment for the deployment (JSON string)",
     )
+    parser.add_argument(
+        "-trc",
+        "--trust_remote_code",
+        default=False,
+        action="store_true",
+        help="Pass trust_remote_code=True to load_tokenizer() for checkpoints whose tokenizer requires it.",
+    )
     return parser.parse_args()
 
 
@@ -283,6 +290,7 @@ def main():
         micro_batch_size=args.micro_batch_size,
         batch_wait_timeout_s=args.batch_wait_timeout_s,
         legacy_model_format=args.legacy_model_format,
+        trust_remote_code=args.trust_remote_code,
         **model_config_kwargs,
     )
 

From d2112a73d56b7a1fe8f7d934b6559e62257557cb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?oliver=20k=C3=B6nig?= <okoenig@nvidia.com>
Date: Wed, 1 Jul 2026 16:13:46 +0200
Subject: [PATCH 5/5] feat: address new inference updates in mcore (#711)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: oliver könig <okoenig@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index c0d5666d4..842ff99d1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -123,6 +123,7 @@ no-build-isolation-package = [
     "mamba-ssm",
     "causal-conv1d",
     "nv-grouped-gemm",
+    "fast-hadamard-transform",
 ]
 
 # Always apply the build group since dependencies like TE/mcore/nemo-run require build dependencies