megatron-core 0.16.0rc0.dev101399__tar.gz → 0.16.0rc0.dev101484__tar.gz

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: megatron-core
-Version: 0.16.0rc0.dev101399
+Version: 0.16.0rc0.dev101484
 Summary: Megatron Core - a library for efficient and scalable training of transformer based models
 Author-email: NVIDIA <nemo-toolkit@nvidia.com>
 Maintainer-email: NVIDIA <nemo-toolkit@nvidia.com>

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/megatron/core/fusions/fused_indices_converter.py

@@ -6,7 +6,7 @@ from unittest.mock import MagicMock
 import torch
 from packaging import version
 
-from megatron.core.utils import null_decorator
+from megatron.core.utils import experimental_fn, null_decorator
 
 try:
     import triton
@@ -279,6 +279,7 @@ class IndicesToMultihot(torch.autograd.Function):
         return None, grad_probs_indices, None, None
 
 
+@experimental_fn(introduced_with_version='0.11.0rc0')
 def fused_indices_to_multihot(indices, probs_indices, num_of_local_experts):
     """Convert moe topk indices to multihot representation.
 

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/megatron/core/fusions/fused_mla_yarn_rope_apply.py

@@ -6,7 +6,7 @@ from unittest.mock import MagicMock
 import torch
 from packaging import version
 
-from megatron.core.utils import null_decorator
+from megatron.core.utils import experimental_fn, null_decorator
 
 try:
     import triton
@@ -324,6 +324,7 @@ class ApplyMLARotaryEmbQ(torch.autograd.Function):
         return grad, None, None, None, None, None, None, None, None
 
 
+@experimental_fn(introduced_with_version="0.13.0")
 def fused_apply_mla_rope_for_q(
     t: torch.Tensor,
     cos: torch.Tensor,
@@ -732,6 +733,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function):
         return d_kv, d_emb, None, None, None, None, None, None, None, None, None
 
 
+@experimental_fn(introduced_with_version="0.13.0")
 def fused_apply_mla_rope_for_kv(
     kv: torch.Tensor,
     k_pos_emb: torch.Tensor,

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/megatron/core/fusions/fused_pad_routing_map.py

@@ -6,7 +6,7 @@ import torch
 from packaging import version
 
 from megatron.core.jit import jit_fuser
-from megatron.core.utils import null_decorator
+from megatron.core.utils import experimental_fn, null_decorator
 
 try:
     import triton
@@ -70,6 +70,7 @@ def _pad_routing_map_kernel(
     tl.store(output_row_ptr + token_indices, output_row, mask=token_mask)
 
 
+@experimental_fn(introduced_with_version="0.13.0")
 @jit_fuser
 def fused_pad_routing_map(routing_map: torch.Tensor, pad_multiple: int) -> torch.Tensor:
     """Fused version of pad_routing_map.

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/megatron/core/inference/contexts/dynamic_block_allocator.py

@@ -115,9 +115,4 @@ class BlockAllocator:
         This resets the available block count to the entire memory pool
         (except for the dummy block).
         """
-        # Reset block bag to so we start consuming
-        # from the beginning of the pool for UVM performance
-        self.block_bag = torch.arange(
-            self.total_count, dtype=torch.int32, device=torch.cuda.current_device()
-        )
         self.total_avail = self.total_count - 1

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/megatron/core/inference/contexts/dynamic_context.py

@@ -578,18 +578,19 @@ class DynamicInferenceContext(BaseInferenceContext):
             """Allocate the memory buffer. This function is called below within
             `with ctx_manager:`."""
             if self.cache_mla_latent:
-                self.memory_buffer = torch.empty(
+                self.memory_buffer = torch.full(
                     (
                         self.num_attention_layers,
                         self.block_allocator.total_count,
                         self.block_size_tokens,
                         self.kv_reduced_dim,
                     ),
+                    -1,
                     dtype=self.params_dtype,
                     device=torch.cuda.current_device(),
                 )
             else:
-                self.memory_buffer = torch.empty(
+                self.memory_buffer = torch.full(
                     (
                         2,  # key and value
                         self.num_attention_layers,
@@ -598,6 +599,7 @@ class DynamicInferenceContext(BaseInferenceContext):
                         self.num_attention_heads_per_partition,
                         self.hidden_size_per_attention_head,
                     ),
+                    -1,
                     dtype=self.params_dtype,
                     device=torch.cuda.current_device(),
                 )
@@ -608,12 +610,12 @@ class DynamicInferenceContext(BaseInferenceContext):
             `with ctx_manager:`."""
             if self.is_hybrid_model:
                 self.mamba_metadata = MambaMetadata(max_requests=self.max_total_requests)
-                self.mamba_conv_states = torch.empty(
+                self.mamba_conv_states = torch.zeros(
                     (self.num_mamba_layers, self.max_total_requests) + self.mamba_conv_states_shape,
                     dtype=self.params_dtype,
                     device=torch.cuda.current_device(),
                 )
-                self.mamba_ssm_states = torch.empty(
+                self.mamba_ssm_states = torch.zeros(
                     (self.num_mamba_layers, self.max_total_requests) + self.mamba_ssm_states_shape,
                     dtype=self.params_dtype,
                     device=torch.cuda.current_device(),
@@ -998,6 +1000,8 @@ class DynamicInferenceContext(BaseInferenceContext):
     def reset_mamba_state(self) -> None:
         """Reset state used within Mamba layers."""
         if self.is_hybrid_model:
+            self.mamba_conv_states.fill_(0)
+            self.mamba_ssm_states.fill_(0)
             self.mamba_metadata.reset()
 
     def using_cuda_graph_this_step(self) -> bool:

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/megatron/core/inference/data_parallel_inference_coordinator.py

@@ -109,8 +109,6 @@ class DataParallelInferenceCoordinator:
             self.identities_of_data_parallel_ranks.append(identity)
         logging.info("Inference Coordinator: Connected with data parallel ranks...")
         self.data_parallel_rank_iterator = cycle(self.identities_of_data_parallel_ranks)
-        self.data_parallel_pause_acks = set()
-        self.data_parallel_stop_acks = set()
 
         self.request_id_to_client_id = {}
         self.request_id_to_client_request_id = {}
@@ -153,7 +151,7 @@ class DataParallelInferenceCoordinator:
                 # print(f"New client connected: {sender_identity}")
                 known_clients.add(sender_identity)
                 self.router_socket.send_multipart(
-                    [sender_identity, msgpack.packb([Headers.CONNECT_ACK.value], use_bin_type=True)]
+                    [sender_identity, msgpack.packb([Headers.ACK.value], use_bin_type=True)]
                 )
 
             elif header == Headers.SUBMIT_REQUEST:
@@ -210,50 +208,6 @@ class DataParallelInferenceCoordinator:
                     self.router_socket.send_multipart(
                         [data_parallel_rank_id, msgpack.packb([header.value], use_bin_type=True)]
                     )
-                if header == Headers.UNPAUSE:
-                    self.data_parallel_pause_acks = set()
-            elif header == Headers.PAUSE_ACK:
-                # control signal ack from the engine
-                assert sender_identity in self.identities_of_data_parallel_ranks
-                assert sender_identity not in self.data_parallel_pause_acks
-                self.data_parallel_pause_acks.add(sender_identity)
-                # route to all clients only once we have gotten an ack from all data parallel ranks
-                if len(self.data_parallel_pause_acks) == self.data_parallel_size:
-                    for client_id in known_clients:
-                        self.router_socket.send_multipart(
-                            [
-                                client_id,
-                                msgpack.packb([header.value, sender_identity], use_bin_type=True),
-                            ]
-                        )
-                    for data_parallel_rank_id in self.identities_of_data_parallel_ranks:
-                        self.router_socket.send_multipart(
-                            [
-                                data_parallel_rank_id,
-                                msgpack.packb([Headers.PAUSE_ACK.value], use_bin_type=True),
-                            ]
-                        )
-            elif header == Headers.STOP_ACK:
-                # control signal ack from the engine
-                assert sender_identity in self.identities_of_data_parallel_ranks
-                assert sender_identity not in self.data_parallel_stop_acks
-                self.data_parallel_stop_acks.add(sender_identity)
-                # route to all clients only once we have gotten an ack from all data parallel ranks
-                if len(self.data_parallel_stop_acks) == self.data_parallel_size:
-                    for client_id in known_clients:
-                        self.router_socket.send_multipart(
-                            [
-                                client_id,
-                                msgpack.packb([header.value, sender_identity], use_bin_type=True),
-                            ]
-                        )
-                    for data_parallel_rank_id in self.identities_of_data_parallel_ranks:
-                        self.router_socket.send_multipart(
-                            [
-                                data_parallel_rank_id,
-                                msgpack.packb([Headers.STOP_ACK.value], use_bin_type=True),
-                            ]
-                        )
             elif header == Headers.ENGINE_REPLY:
                 # This is the output of a single engine step on some data parallel rank.
                 assert sender_identity in self.identities_of_data_parallel_ranks
@@ -270,7 +224,7 @@ class DataParallelInferenceCoordinator:
                         [
                             client_identity,
                             msgpack.packb(
-                                [header.value, client_request_identity, finished_request_record],
+                                [client_request_identity, finished_request_record],
                                 use_bin_type=True,
                             ),
                         ]

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/megatron/core/inference/engines/dynamic_engine.py

@@ -168,6 +168,7 @@ class DynamicInferenceEngine(AbstractEngine):
         self.enable_chunked_prefill = enable_chunked_prefill
         self.inference_logging_step_interval = inference_logging_step_interval
         self.unified_memory_level = context.unified_memory_level
+        self.capture_stats = None  # pass CI
 
         if enable_cuda_graph is not None:
             self.cuda_graph_impl = "local" if enable_cuda_graph else "none"
@@ -228,11 +229,8 @@ class DynamicInferenceEngine(AbstractEngine):
         # Runtime state.
         self._loop = get_asyncio_loop(getattr(self, "_loop", None))
         self._cond = asyncio.Condition()
-        self.running = asyncio.Event()
-        self.paused = asyncio.Event()
-        self.stopped = asyncio.Event()
-        self.received_pause: bool = False
-        self.received_stop: bool = False
+        self.paused = False
+        self.stopped = False
         self.suspend_signal = False
         self.is_suspended = False
         self.resume_request_ids = None
@@ -1316,41 +1314,14 @@ class DynamicInferenceEngine(AbstractEngine):
         for message in all_messages:
             data = msgpack.unpackb(message, raw=False)
             header = Headers(data[0])
-
-            if self.received_stop:
-                assert (
-                    header == Headers.STOP_ACK
-                ), "Engine is shutting down. No other messages allowed except STOP_ACK."
-
             if header == Headers.SUBMIT_REQUEST:
                 request_id, prompt, sampling_params = data[1:]
                 sampling_params = SamplingParams.deserialize(sampling_params)
                 self.add_request(request_id, prompt, sampling_params)
             elif header == Headers.PAUSE:
-                # Pause thyself.
-                self.received_pause = True
-                self.running.clear()
-                # Send PAUSE_ACK back to coordinator.
-                if self.is_mp_coordinator:
-                    payload = msgpack.packb([Headers.PAUSE_ACK.value], use_bin_type=True)
-                    self.socket_for_receiving_requests.send(payload)
-            elif header == Headers.STOP:
-                # Stop thyself.
-                self.received_stop = True
-                self.running.clear()
-                # Send STOP_ACK back to coordinator.
-                if self.is_mp_coordinator:
-                    payload = msgpack.packb([Headers.STOP_ACK.value], use_bin_type=True)
-                    self.socket_for_receiving_requests.send(payload)
-            elif header == Headers.PAUSE_ACK:
-                self.paused.set()
-                self.received_pause = False
-            elif header == Headers.STOP_ACK:
-                self.stopped.set()
-                self.stop()
+                self.paused = True
             elif header == Headers.UNPAUSE:
-                self.paused.clear()
-                self.running.set()
+                self.paused = False
             elif header == Headers.SUSPEND:
                 self.suspend_signal = True
             elif header == Headers.RESUME:
@@ -1411,10 +1382,11 @@ class DynamicInferenceEngine(AbstractEngine):
         try:
             while True:
                 self.schedule_requests()
-                if self.stopped.is_set():
-                    break
+                if self.stopped:
+                    self.stop()
+                    return
 
-                # for the cases below (no active requests, or undergoing a state-change)
+                # for the cases below (engine is paused or no active requests),
                 # do not use asyncio.sleep(0)
                 # as tp-rank=0 will flood the num_messages publisher
                 # with "0" repeatedly. This causes some packets to drop.
@@ -1426,7 +1398,7 @@ class DynamicInferenceEngine(AbstractEngine):
 
                 # todo [Siddharth]: Can this hardcoded sleep be avoided
                 # with asyncio zmq sockets?
-                if self.paused.is_set() or self.received_pause or self.received_stop:
+                if self.paused:
                     await asyncio.sleep(0.02)
                     continue
 

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/megatron/core/inference/headers.py

@@ -9,16 +9,14 @@ class Headers(Enum):
     """
 
     CONNECT = auto()
-    CONNECT_ACK = auto()
+    ACK = auto()
     SUBMIT_REQUEST = auto()
     ENGINE_REPLY = auto()
     PAUSE = auto()
-    PAUSE_ACK = auto()
     UNPAUSE = auto()
     SUSPEND = auto()
     RESUME = auto()
     STOP = auto()
-    STOP_ACK = auto()
 
 
 class UnknownHeaderError(Exception):

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/megatron/core/inference/inference_client.py

@@ -4,7 +4,7 @@ import asyncio
 import logging
 import os
 import time
-from typing import Awaitable, List, Optional, Union
+from typing import List, Union
 
 from megatron.core.inference.inference_request import DynamicInferenceRequestRecord
 from megatron.core.inference.sampling_params import SamplingParams
@@ -73,11 +73,6 @@ class InferenceClient:
         inference_coordinator_address = os.getenv('MASTER_ADDR', '127.0.0.1')
         socket.connect(f"tcp://{inference_coordinator_address}:{inference_coordinator_port}")
 
-        self._loop = None
-        self.running = asyncio.Event()
-        self.paused = asyncio.Event()
-        self.stopped = asyncio.Event()
-
         self.socket = socket
         self.completion_futures = {}
         self.request_submission_times = {}
@@ -103,49 +98,35 @@ class InferenceClient:
             asyncio.Future: A future that will be resolved with a
             `DynamicInferenceRequestRecord` object containing the completed result.
         """
-        if not self.running.is_set():
-            raise RuntimeError("InferenceClient is not currently running.")
         request_id = self.next_request_id
         self.next_request_id += 1
         payload = [Headers.SUBMIT_REQUEST.value, request_id, prompt, sampling_params.serialize()]
         payload_serialized = msgpack.packb(payload, use_bin_type=True)
         self.socket.send(payload_serialized)
         assert request_id not in self.completion_futures
-        self.completion_futures[request_id] = self._loop.create_future()
+        self.completion_futures[request_id] = get_asyncio_loop().create_future()
         self.request_submission_times[request_id] = time.perf_counter()
         return self.completion_futures[request_id]
 
     @trace_async_exceptions
-    async def _recv_task(self):
+    async def _listen_for_completed_requests(self):
         """
         Listens for completed inference requests from the coordinator.
 
         This coroutine runs in an infinite loop, continuously polling the socket
-        for data.
-        When a request reply is received, it unpacks the message, finds the
+        for replies. When a reply is received, it unpacks the message, finds the
         corresponding Future using the request ID, and sets the result.
-        Other control packets are handled appropriately.
 
         This method is started as a background task by the `start()` method.
         """
         while True:
             try:
-                data = msgpack.unpackb(self.socket.recv(flags=zmq.NOBLOCK), raw=False)
-                header = Headers(data[0])
-                if header == Headers.ENGINE_REPLY:
-                    request_id, reply = data[1:]
-                    reply['latency'] = time.perf_counter() - self.request_submission_times.pop(
-                        request_id
-                    )
-                    completion_future = self.completion_futures.pop(request_id)
-                    if completion_future.done():
-                        logging.warning(f"Client: The future for {request_id} has been cancelled!")
-                        continue
-                    completion_future.set_result(DynamicInferenceRequestRecord.deserialize(reply))
-                elif header == Headers.PAUSE_ACK:
-                    self.paused.set()
-                elif header == Headers.STOP_ACK:
-                    self.stopped.set()
+                request_id, reply = msgpack.unpackb(self.socket.recv(flags=zmq.NOBLOCK), raw=False)
+                reply['latency'] = time.perf_counter() - self.request_submission_times.pop(
+                    request_id
+                )
+                completion_future = self.completion_futures.pop(request_id)
+                completion_future.set_result(DynamicInferenceRequestRecord.deserialize(reply))
             except zmq.Again:
                 await asyncio.sleep(0.005)
                 continue
@@ -156,15 +137,15 @@ class InferenceClient:
         """
         Performs the initial handshake with the inference coordinator.
 
-        Sends a CONNECT signal and waits for a CONNECT_ACK reply to ensure the
+        Sends a CONNECT signal and waits for an ACK reply to ensure the
         connection is established and acknowledged by the coordinator.
         """
         payload = [Headers.CONNECT.value]
         self.socket.send(msgpack.packb(payload, use_bin_type=True))
         reply = msgpack.unpackb(self.socket.recv(), raw=False)[0]
-        assert Headers(reply) == Headers.CONNECT_ACK
+        assert Headers(reply) == Headers.ACK
 
-    async def start(self, loop: Optional[asyncio.AbstractEventLoop] = None):
+    async def start(self):
         """
         Connects to the coordinator and starts the background listener task.
 
@@ -173,12 +154,8 @@ class InferenceClient:
         coroutine.
         """
         logging.info("Client: Connecting to InferenceCoordinator...")
-        self._loop = get_asyncio_loop(loop)
-        self.running.set()
-        self.paused.clear()
-        self.stopped.clear()
         self._connect_with_inference_coordinator()
-        self.listener_task = self._loop.create_task(self._recv_task())
+        self.listener_task = asyncio.create_task(self._listen_for_completed_requests())
 
     def _send_signal_to_engines(self, signal):
         """
@@ -191,25 +168,12 @@ class InferenceClient:
         payload_serialized = msgpack.packb(payload, use_bin_type=True)
         self.socket.send(payload_serialized)
 
-    def pause_engines(self) -> Awaitable:
-        """Sends a signal to pause all inference engines.
-
-        The signal first propagates thru the coordinator to all engines.
-        All engines acknowledge this signal and clear their `running` flags.
-        The coordinator awaits all acknowledgements before forwarding the ACK
-            back to the client, as well as to the engines.
-        The engines set their `paused` flags upon seeing the ACK.
-
-        Returns:
-            Awaitable: An awaitable that resolves when all engines have paused.
-        """
+    def pause_engines(self):
+        """Sends a signal to pause all inference engines."""
         self._send_signal_to_engines(Headers.PAUSE)
-        return self.paused.wait()
 
-    def unpause_engines(self) -> None:
+    def unpause_engines(self):
         """Sends a signal to unpause all inference engines."""
-        self.paused.clear()
-        self.running.set()
         self._send_signal_to_engines(Headers.UNPAUSE)
 
     def suspend_engines(self):
@@ -222,21 +186,9 @@ class InferenceClient:
         self._send_signal_to_engines(Headers.RESUME)
         self._send_signal_to_engines(Headers.UNPAUSE)
 
-    def stop_engines(self) -> Awaitable:
-        """Sends a signal to gracefully stop all inference engines.
-
-        The signal first propagates thru the coordinator to all engines.
-        All engines acknowledge this signal and clear their `running` flags.
-        The coordinator awaits all acknowledgements before forwarding the ACK
-            back to the client, as well as to the engines.
-        The engines set their `stopped` flags upon seeing the ACK.
-
-        Returns:
-            Awaitable: An awaitable that resolves when all engines have stopped.
-        """
+    def stop_engines(self):
+        """Sends a signal to gracefully stop all inference engines."""
         self._send_signal_to_engines(Headers.STOP)
-        self.running.clear()
-        return self.stopped.wait()
 
     def stop(self):
         """

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/megatron/core/inference/unified_memory.py

@@ -1,9 +1,7 @@
 # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
 
 import os
-import signal
 import warnings
-from contextlib import contextmanager
 from enum import Enum, auto
 from pathlib import Path
 
@@ -30,45 +28,16 @@ class CompilationState(Enum):
     SUCCESS = auto()  # Compilation attempted, and succeeded.
 
 
-class UnifiedMemoryUnsupportedError(Exception):
-    """Unified memory is not supported on this system."""
-
-
-class UnifiedMemoryCompileTimeoutError(UnifiedMemoryUnsupportedError):
-    """Unified memory compilation timed out."""
-
-
 # Compilation vars.
 _compilation_state = CompilationState.UNATTEMPTED
 _alloc = None  # must remain global until process exit.
 _mod = None  # must remain global until process exit.
 
 
-@contextmanager
-def _compile_timeout(timeout_s: int):
-    """Context manager to timeout compilation.
-
-    Args:
-        timeout_s (int): Timeout in seconds.
-    """
-
-    def _handler(signum, frame):
-        raise UnifiedMemoryCompileTimeoutError(
-            "Unified memory compilation has been forcefully timed out. "
-            "This is almost certainly due to stale lock files associated with your Unix user. "
-            "The official PyTorch advice is to resolve this issue with the following command:\n"
-            "`rm -rf /tmp/torch_extensions/`\n"
-            "Alternately, the TORCH_EXTENSIONS_DIR env var may be set to a different path. "
-            "Please clean up your stale cache and try again."
-        )
+class UnifiedMemoryUnsupportedError(Exception):
+    """Unified memory is not supported on this system."""
 
-    curr_handler = signal.signal(signal.SIGALRM, _handler)
-    try:
-        signal.alarm(timeout_s)
-        yield
-    finally:
-        signal.alarm(0)
-        signal.signal(signal.SIGALRM, curr_handler)
+    pass
 
 
 def compile_allocator():
@@ -136,9 +105,6 @@ def compile_allocator():
     }
     """
 
-    # Define a timeout of 30s for how long the build is allowed to run.
-    timeout_s = 30
-
     # Build the .so upon import; this avoids issues.
     if _has_mem_pool:
         _extra_ldflags = ["-lcudart"]
@@ -147,20 +113,18 @@ def compile_allocator():
             if os.path.isdir(_cuda_lib):
                 _extra_ldflags = [f"-L{_cuda_lib}", "-lcudart"]
         try:
-            with _compile_timeout(timeout_s):
-                _mod = load_inline(
-                    name="managed_alloc_runtime",
-                    cpp_sources=[_mempool_c_src],
-                    functions=[],
-                    with_cuda=True,
-                    extra_ldflags=_extra_ldflags,
-                    verbose=True,
-                )
-                _so_path = Path(_mod.__file__).as_posix()
-                _cpa = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free")
-                _alloc = _cpa.allocator()
-                _compilation_state = CompilationState.SUCCESS
-        except (RuntimeError, ImportError, OSError, UnifiedMemoryCompileTimeoutError) as e:
+            _mod = load_inline(
+                name="managed_alloc_runtime",
+                cpp_sources=[_mempool_c_src],
+                functions=[],
+                with_cuda=True,
+                extra_ldflags=_extra_ldflags,
+                verbose=True,
+            )
+            _so_path = Path(_mod.__file__).as_posix()
+            _alloc = CUDAPluggableAllocator(_so_path, "managed_malloc", "managed_free").allocator()
+            _compilation_state = CompilationState.SUCCESS
+        except (RuntimeError, ImportError, OSError) as e:
             warnings.warn(f"Failed to create unified memory mempool: '{e}'.")
             _compilation_state = CompilationState.FAILURE
 

{megatron_core-0.16.0rc0.dev101399 → megatron_core-0.16.0rc0.dev101484}/megatron/core/models/common/embeddings/rope_utils.py

@@ -268,21 +268,20 @@ def apply_rotary_pos_emb(
     if config.apply_rope_fusion:
         if cu_seqlens is None:
             # NOTE: TE backends do not support mRoPE in bshd format when bs > 1.
-            use_unfused = False
             if config.mrope_section is not None and freqs.shape[1] > 1:
                 # TODO: Add a check in TransformerConfig and remove this unfused implementation.
                 warnings.warn(
                     "apply_rope_fusion does not support mRoPE in bshd format when bs > 1. "
                     "Please set apply_rope_fusion to false. This will become an error in v0.16."
                 )
-                use_unfused = True
-            if mscale != 1.0:
-                warnings.warn(
-                    f"mscale={mscale} is not supported by TE's fused RoPE. "
-                    "Using unfused implementation."
+                return _apply_rotary_pos_emb_bshd(
+                    t,
+                    freqs,
+                    rotary_interleaved=config.rotary_interleaved,
+                    multi_latent_attention=config.multi_latent_attention,
+                    mscale=mscale,
                 )
-                use_unfused = True
-            if not use_unfused:
+            else:
                 assert fused_apply_rotary_pos_emb is not None, "apply_rope_fusion is not available."
                 return fused_apply_rotary_pos_emb(t, freqs, interleaved=config.rotary_interleaved)
         else:
@@ -290,25 +289,25 @@ def apply_rotary_pos_emb(
             return fused_apply_rotary_pos_emb_thd(
                 t, cu_seqlens, freqs, cp_size=cp_group.size(), cp_rank=cp_group.rank()
             )
-    # use unfused implementation
-    if cu_seqlens is None:
-        return _apply_rotary_pos_emb_bshd(
-            t,
-            freqs,
-            rotary_interleaved=config.rotary_interleaved,
-            multi_latent_attention=config.multi_latent_attention,
-            mscale=mscale,
-        )
     else:
-        return _apply_rotary_pos_emb_thd(
-            t,
-            cu_seqlens,
-            freqs,
-            rotary_interleaved=config.rotary_interleaved,
-            multi_latent_attention=config.multi_latent_attention,
-            mscale=mscale,
-            cp_group=cp_group,
-        )
+        if cu_seqlens is None:
+            return _apply_rotary_pos_emb_bshd(
+                t,
+                freqs,
+                rotary_interleaved=config.rotary_interleaved,
+                multi_latent_attention=config.multi_latent_attention,
+                mscale=mscale,
+            )
+        else:
+            return _apply_rotary_pos_emb_thd(
+                t,
+                cu_seqlens,
+                freqs,
+                rotary_interleaved=config.rotary_interleaved,
+                multi_latent_attention=config.multi_latent_attention,
+                mscale=mscale,
+                cp_group=cp_group,
+            )
 
 
 def apply_rotary_pos_emb_with_cos_sin(