Remove d2d page feature for now

Remove unused triton kernel for dropping token in case overflow happens

Author: nanz-nv

megatron/core/transformer/moe/moe_utils.py

@@ -34,11 +34,6 @@
     HAVE_TE = False
 
 
-import triton
-import triton.language as tl
-
-
-
 # MOE logging
 _MOE_LAYER_WISE_LOGGING_TRACKER = {}
 
@@ -934,7 +929,7 @@ def forward(ctx, logits):
         """
         Forward pass returns random logits with rank-specific seed.
         """
-        if RandomSTE.random_logits is not None:
+        if is_graph_capturing() and RandomSTE.random_logits is not None:
             return RandomSTE.random_logits
 
         if RandomSTE.generator is None:
@@ -1302,95 +1297,3 @@ def wrapped_func(moe_layer, *args, **kwargs):
         return wrapped_func
 
     return decorator
-
-@triton.jit
-def _drop_routing_map_kernel(
-    routing_map_ptr,
-    over_budget_ptr,
-    routing_map_dropped_ptr,
-    num_elements: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    """Triton kernel to drop routing map based on budget constraints.
-    
-    Args:
-        routing_map_ptr: Pointer to the input routing_map tensor
-        over_budget_ptr: Pointer to the boolean tensor indicating if any EP rank is over budget
-        routing_map_dropped_ptr: Pointer to the output routing_map tensor
-        num_elements: Total number of elements to process
-        BLOCK_SIZE: Block size for Triton kernel
-    """
-    # Get the program ID
-    pid = tl.program_id(axis=0)
-    
-    # Read the over_budget value (scalar tensor with single element)
-    over_budget_val = tl.load(over_budget_ptr)
-    
-    # Calculate the offset for this program
-    offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    
-    # Load the routing_map values
-    mask = offset < num_elements
-    routing_map_val = tl.load(routing_map_ptr + offset, mask=mask, other=0.0)
-    
-    # If over_budget is 1 (True), output is 0 (drop); if over_budget is 0 (False), output is routing_map_val (keep)
-    output_val = routing_map_val * (1 - over_budget_val)
-    
-    # Store the result
-    tl.store(routing_map_dropped_ptr + offset, output_val, mask=mask)
-
-
-def drop_routing_map_triton(
-    routing_map: torch.Tensor, 
-    budget: torch.Tensor, 
-    num_tokens_per_ep_rank: torch.Tensor
-) -> torch.Tensor:
-    """Drop tokens from routing_map that exceed the budget per EP rank using Triton.
-    
-    Args:
-        routing_map: Tensor indicating which tokens are assigned to each expert.
-        budget: Integer tensor with the maximum number of tokens per EP rank.
-        num_tokens_per_ep_rank: Tensor with actual number of tokens per EP rank.
-    
-    Returns:
-        Modified routing_map with tokens exceeding budget zeroed out if any EP rank 
-        exceeds budget, otherwise returns the original routing_map.
-    """
-    
-    # Calculate boolean tensor: over_budget is True if ANY EP rank exceeds budget
-    over_budget = (num_tokens_per_ep_rank > budget).any()
-    
-    # Convert boolean to int8 
-    over_budget_int = over_budget.to(torch.int8)
-    
-    # Convert routing_map to numeric type if it's boolean
-    if routing_map.dtype == torch.bool:
-        routing_map_numeric = routing_map.to(torch.int8)
-    else:
-        routing_map_numeric = routing_map
-    
-    # Create output tensor with same dtype as input
-    routing_map_dropped = torch.empty_like(routing_map_numeric)
-    
-    # Flatten tensors for kernel processing
-    routing_map_flat = routing_map_numeric.flatten()
-    num_elements = routing_map_flat.numel()
-    
-    # Determine grid size
-    BLOCK_SIZE = 1024
-    grid = (triton.cdiv(num_elements, BLOCK_SIZE),)
-    
-    # Launch kernel with over_budget tensor pointer (as int8)
-    _drop_routing_map_kernel[grid](
-        routing_map_flat,
-        over_budget_int,
-        routing_map_dropped.flatten(),
-        num_elements,
-        BLOCK_SIZE=BLOCK_SIZE,
-    )
-    
-    # Convert back to boolean if original was boolean
-    if routing_map.dtype == torch.bool:
-        routing_map_dropped = routing_map_dropped.to(torch.bool)
-    
-    return routing_map_dropped, over_budget.to(torch.bool)

megatron/core/transformer/moe/paged_stash.py

@@ -247,7 +247,7 @@ class PagedTensor:
     A paged tensor that stores data in pages within a paged stash buffer.
     """
 
-    def __init__(self, tensor, num_tokens_tensor=None, vp_stage=None, schedule_layer_no=None, layer_name=None, max_tokens=None, page_size=64, num_d2d_pages=0):
+    def __init__(self, tensor, num_tokens_tensor=None, vp_stage=None, schedule_layer_no=None, layer_name=None, max_tokens=None, page_size=64):
         """
         Args:
             tensor: The tensor to store
@@ -256,7 +256,6 @@ def __init__(self, tensor, num_tokens_tensor=None, vp_stage=None, schedule_layer
             layer_name: Name of the layer
             max_tokens: Maximum number of tokens
             page_size: Number of tokens per page
-            num_d2d_pages: Number of pages to copy using native PyTorch (rest uses Triton)
         """
         self._tensor = tensor
         self._original_tensor = None
@@ -267,7 +266,6 @@ def __init__(self, tensor, num_tokens_tensor=None, vp_stage=None, schedule_layer
         self.layer_name = layer_name
         self.max_tokens = max_tokens
         self.page_size = page_size
-        self.num_d2d_pages = num_d2d_pages
 
         # Original tensor information
         self.original_shape = list(tensor.shape)
@@ -282,13 +280,6 @@ def __init__(self, tensor, num_tokens_tensor=None, vp_stage=None, schedule_layer
 
         # Page record: stores which pages are being used for this tensor
         self.page_record = torch.zeros(self.max_num_pages, dtype=torch.int64, device=self.device)
-        
-        # Static tensor for D2D pages (allocate upfront if needed)
-        d2d_tokens = min(self.num_d2d_pages * self.page_size, self.max_num_tokens)
-        if d2d_tokens > 0:
-            self.static_tensor = torch.empty((d2d_tokens, self.hidden_size), dtype=self.dtype, device=self.device)
-        else:
-            self.static_tensor = None
 
     @property
     def schedule_layer(self):
@@ -312,48 +303,33 @@ def offload_to_stash(self, paged_stash_buffer: PagedStashBuffer, max_blocks=2048
         else:
             tensor_to_copy = self._tensor
 
-        # Split tensor into two parts: D2D portion and Triton portion
-        # Use max_num_tokens for consistent size across iterations
-        d2d_tokens = min(self.num_d2d_pages * self.page_size, self.max_num_tokens)
-        triton_tokens = self.max_num_tokens - d2d_tokens
-        
-        # Perform both D2D copy and Triton kernel together
-        # Part 1: Copy first d2d_tokens to static_tensor using native PyTorch
-        if d2d_tokens > 0:
-            self.static_tensor[:d2d_tokens] = tensor_to_copy[:d2d_tokens]
-        # Part 2: Copy remaining tokens using Triton kernel
-        if triton_tokens > 0:
-            triton_tensor = tensor_to_copy[d2d_tokens:self.max_num_tokens]
-            # Use actual num_tokens for the kernel (how many tokens to actually copy)
-            triton_num_tokens = self.num_tokens_tensor - d2d_tokens
-            
-            # Determine grid size
-            BLOCK_SIZE = GLOBAL_BLOCK_SIZE
-            num_blocks = min(triton_tokens, max_blocks)
-            grid = (num_blocks,)
-            
-            # Create temporary tensor for new head
-            new_free_list_head = torch.empty(1, dtype=torch.int64, device=self.device)
-            
-            # Launch paged stash copy kernel
-            _paged_stash_copy_kernel[grid](
-                triton_tensor,
-                paged_stash_buffer.buffer,
-                triton_num_tokens,
-                paged_stash_buffer.free_list,
-                paged_stash_buffer.free_list_head,
-                paged_stash_buffer.free_list_tail,
-                paged_stash_buffer.free_list_capacity,
-                self.page_record,  # Triton kernel will populate page_record
-                paged_stash_buffer.overflow,
-                new_free_list_head,
-                PAGE_SIZE=self.page_size,
-                HIDDEN_SIZE=self.hidden_size,
-                BLOCK_SIZE=BLOCK_SIZE,
-            )
-            
-            # Update free list head
-            paged_stash_buffer.free_list_head.copy_(new_free_list_head)
+        # Determine grid size
+        BLOCK_SIZE = GLOBAL_BLOCK_SIZE
+        num_blocks = min(self.max_num_tokens, max_blocks)
+        grid = (num_blocks,)
+        
+        # Create temporary tensor for new head
+        new_free_list_head = torch.empty(1, dtype=torch.int64, device=self.device)
+        
+        # Launch paged stash copy kernel
+        _paged_stash_copy_kernel[grid](
+            tensor_to_copy,
+            paged_stash_buffer.buffer,
+            self.num_tokens_tensor,
+            paged_stash_buffer.free_list,
+            paged_stash_buffer.free_list_head,
+            paged_stash_buffer.free_list_tail,
+            paged_stash_buffer.free_list_capacity,
+            self.page_record,  # Triton kernel will populate page_record
+            paged_stash_buffer.overflow,
+            new_free_list_head,
+            PAGE_SIZE=self.page_size,
+            HIDDEN_SIZE=self.hidden_size,
+            BLOCK_SIZE=BLOCK_SIZE,
+        )
+        
+        # Update free list head
+        paged_stash_buffer.free_list_head.copy_(new_free_list_head)
 
         # Save reference to original tensor
         self._original_tensor = self._tensor
@@ -384,48 +360,32 @@ def reload_from_stash(self, paged_stash_buffer: PagedStashBuffer, max_blocks=204
             self._tensor = torch.empty(self.original_shape, dtype=self.dtype, device=self.device)
             tensor_to_reload = self._tensor
 
-        # Split tensor into two parts: D2D portion and Triton portion
-        # Use max_num_tokens for consistency with stash
-        d2d_tokens = min(self.num_d2d_pages * self.page_size, self.max_num_tokens)
-        triton_tokens = self.max_num_tokens - d2d_tokens
-        
-        # Perform both D2D copy and Triton kernel together
-        # Part 1: Copy first d2d_tokens from static_tensor using native PyTorch
-        if d2d_tokens > 0 and self.static_tensor is not None:
-            tensor_to_reload[:d2d_tokens] = self.static_tensor[:d2d_tokens]
-        
-        # Part 2: Copy remaining tokens using Triton kernel
-        if triton_tokens > 0:
-            triton_tensor = tensor_to_reload[d2d_tokens:self.max_num_tokens]
-            # Use actual num_tokens for the kernel (how many tokens to actually copy)
-            triton_num_tokens = self.num_tokens_tensor - d2d_tokens
-            
-            # Determine grid size
-            BLOCK_SIZE = GLOBAL_BLOCK_SIZE
-            num_blocks = min(triton_tokens, max_blocks)
-            grid = (num_blocks,)
-            
-            # Create temporary tensor for new tail
-            new_free_list_tail = torch.empty(1, dtype=torch.int64, device=self.device)
-            
-            # Launch paged stash pop kernel
-            _paged_stash_pop_kernel[grid](
-                paged_stash_buffer.buffer,
-                triton_tensor,
-                triton_num_tokens,
-                self.page_record,  # Triton kernel will read from page_record
-                paged_stash_buffer.free_list,
-                paged_stash_buffer.free_list_head,
-                paged_stash_buffer.free_list_tail,
-                paged_stash_buffer.free_list_capacity,
-                new_free_list_tail,
-                PAGE_SIZE=self.page_size,
-                HIDDEN_SIZE=self.hidden_size,
-                BLOCK_SIZE=BLOCK_SIZE,
-            )
-            
-            # Update free list tail
-            paged_stash_buffer.free_list_tail.copy_(new_free_list_tail)
+        # Determine grid size
+        BLOCK_SIZE = GLOBAL_BLOCK_SIZE
+        num_blocks = min(self.max_num_tokens, max_blocks)
+        grid = (num_blocks,)
+        
+        # Create temporary tensor for new tail
+        new_free_list_tail = torch.empty(1, dtype=torch.int64, device=self.device)
+        
+        # Launch paged stash pop kernel
+        _paged_stash_pop_kernel[grid](
+            paged_stash_buffer.buffer,
+            tensor_to_reload,
+            self.num_tokens_tensor,
+            self.page_record,  # Triton kernel will read from page_record
+            paged_stash_buffer.free_list,
+            paged_stash_buffer.free_list_head,
+            paged_stash_buffer.free_list_tail,
+            paged_stash_buffer.free_list_capacity,
+            new_free_list_tail,
+            PAGE_SIZE=self.page_size,
+            HIDDEN_SIZE=self.hidden_size,
+            BLOCK_SIZE=BLOCK_SIZE,
+        )
+        
+        # Update free list tail
+        paged_stash_buffer.free_list_tail.copy_(new_free_list_tail)
 
 
 class PP_PreScheduleFunction(torch.autograd.Function):
@@ -555,9 +515,6 @@ def __init__(self):
 
         # Page size for paged memory management
         self.page_size = int(os.getenv('PAGED_STASH_PAGE_SIZE', '64'))  # Default 64 tokens per page
-        
-        # Number of pages to copy using native PyTorch (D2D)
-        self.num_d2d_pages = int(os.getenv('NUM_D2D_PAGES', '0'))  # Default 0 (all Triton)
 
     @property
     def pack_stream(self):
@@ -765,7 +722,6 @@ def on_save_for_backward(self, tensor: torch.Tensor) -> Any:
             layer_name=self._current_layer_name, 
             max_tokens=self.max_num_tokens,
             page_size=self.page_size,
-            num_d2d_pages=self.num_d2d_pages
         )
 
         if self.status == 'captured':

megatron/core/transformer/moe/token_dispatcher.py

@@ -33,7 +33,6 @@
     permute,
     sort_chunks_by_idxs,
     unpermute,
-    drop_routing_map_triton,
 )
 from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
 from megatron.core.transformer.transformer_config import TransformerConfig