feat(moe): allow user-tuned FusedA2A config (chunk_size, num_sms) for DeepEP

docker/Dockerfile.linting

@@ -20,4 +20,4 @@ FROM main AS jet
 ARG JET_API_VERSION
 RUN --mount=type=secret,id=JET_INDEX_URLS \
       JET_INDEX_URLS=$(cat /run/secrets/JET_INDEX_URLS) && \ 
-      uv pip install --no-cache-dir "jet-client~=2.0" --upgrade $JET_INDEX_URLS
+      uv pip install --no-cache-dir "jet-client~=2.0" --upgrade $JET_INDEX_URLS

megatron/core/transformer/moe/fused_a2a.py

@@ -7,16 +7,19 @@
 
 from megatron.core.utils import internal_api
 
+import torch
+
+from megatron.core.transformer.moe.fused_a2a_config import FusedA2AConfig
+
 try:
     from deep_ep import Buffer
     from deep_ep.utils import EventHandle, EventOverlap
+    from deep_ep_cpp import Config as DeepEPConfig
 
     HAVE_DEEP_EP = True
 except ImportError:
     HAVE_DEEP_EP = False
 
-import torch
-
 _buffer = None
 
 
@@ -68,6 +71,56 @@ def get_buffer(group: torch.distributed.ProcessGroup, hidden_bytes: int):
     return _buffer
 
 
+if HAVE_DEEP_EP:
+
+    def _build_deepep_config(default_config, fused_a2a_cfg):
+        """Build a DeepEP Config from a default config and user overrides in FusedA2AConfig.
+
+        Starts from the hardware-appropriate default (e.g. Buffer.get_dispatch_config) and
+        overrides only the fields the user explicitly set.  Returns the default unchanged
+        when no overrides are needed, avoiding any object allocation on the hot path.
+
+        Args:
+            default_config: Config returned by Buffer.get_dispatch_config() or
+                            Buffer.get_combine_config() — hardware-tuned baseline.
+            fused_a2a_cfg:  FusedA2AConfig with optional user overrides, or None.
+
+        Returns:
+            A DeepEP Config (possibly the original default_config if nothing was overridden).
+
+        Note:
+            chunk_size maps to num_max_nvl_chunked_send_tokens.  The DeepEP C++ assertion
+            requires chunk_size < num_max_nvl_chunked_recv_tokens (default 256), so
+            values >= 256 will raise inside the DeepEP C++ constructor.
+        """
+        if fused_a2a_cfg is None:
+            return default_config
+        if fused_a2a_cfg.num_sms is None and fused_a2a_cfg.chunk_size is None:
+            return default_config
+        num_sms = (
+            fused_a2a_cfg.num_sms
+            if fused_a2a_cfg.num_sms is not None
+            else default_config.num_sms
+        )
+        chunk_size = (
+            fused_a2a_cfg.chunk_size
+            if fused_a2a_cfg.chunk_size is not None
+            else default_config.num_max_nvl_chunked_send_tokens
+        )
+        # Build a new Config preserving all RDMA / buffer-size params from the default.
+        # DeepEPConfig == deep_ep_cpp.Config, guaranteed available here (HAVE_DEEP_EP=True).
+        return DeepEPConfig(
+            num_sms,
+            chunk_size,
+            default_config.num_max_nvl_chunked_recv_tokens,
+            default_config.num_max_rdma_chunked_send_tokens,
+            default_config.num_max_rdma_chunked_recv_tokens,
+        )
+
+else:
+    _build_deepep_config = None
+
+
 class FusedDispatch(torch.autograd.Function):
     """Fused dispatch operation for MoE routing combining computation and communication."""
 
@@ -81,11 +134,20 @@ def forward(
         group,
         async_finish=False,
         allocate_on_comm_stream=False,
+        fused_a2a_config=None,
     ):
         """Forward pass of fused dispatch."""
         previous_event = None
         if async_finish:
             previous_event = EventOverlap(EventHandle())
+        # Build hardware-appropriate DeepEP Configs, applying user overrides.
+        # _build_deepep_config reads fields from the C++ Config struct exposed by pybind11.
+        dispatch_config = _build_deepep_config(
+            Buffer.get_dispatch_config(group.size()), fused_a2a_config
+        )
+        combine_config = _build_deepep_config(
+            Buffer.get_combine_config(group.size()), fused_a2a_config
+        )
         # Calculate layout before actual dispatch
         buffer = get_buffer(group, get_hidden_bytes(x))
         (
@@ -120,6 +182,7 @@ def forward(
             num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
             is_token_in_rank=is_token_in_rank,
             num_tokens_per_expert=num_tokens_per_expert,
+            config=dispatch_config,
             previous_event=event,  # wait in deepep::intra/inter_dispatch
             async_finish=async_finish,
             allocate_on_comm_stream=allocate_on_comm_stream,
@@ -129,11 +192,12 @@ def forward(
         if async_finish:
             after_event_overlap.current_stream_wait()
 
-        # Save for backward
+        # Save for backward (combine uses the combine config)
         ctx.group = group
         ctx.handle = handle
         ctx.async_finish = async_finish
         ctx.allocate_on_comm_stream = allocate_on_comm_stream
+        ctx.combine_config = combine_config
         tokens_per_expert = torch.tensor(num_recv_tokens_per_expert_list)
 
         return (recv_x, recv_token_indices, recv_token_probs, tokens_per_expert, handle)
@@ -152,29 +216,41 @@ def backward(
             grad_output.contiguous(),
             handle,
             topk_weights=grad_token_probs.float(),
+            config=ctx.combine_config,
             previous_event=previous_event,
             async_finish=ctx.async_finish,
             allocate_on_comm_stream=ctx.allocate_on_comm_stream,
         )
         # Make sure current stream is synchronized
         if ctx.async_finish:
             after_event.current_stream_wait()
-        return grad_x, None, grad_token_probs, None, None, None, None
+        return grad_x, None, grad_token_probs, None, None, None, None, None
 
 
 class FusedCombine(torch.autograd.Function):
     """Fused combine operation for MoE output combining computation and communication."""
 
     @staticmethod
-    def forward(ctx, x, group, handle, async_finish=False, allocate_on_comm_stream=False):
+    def forward(
+        ctx, x, group, handle, async_finish=False, allocate_on_comm_stream=False,
+        fused_a2a_config=None,
+    ):
         """Forward pass of fused combine."""
         previous_event = None
         if async_finish:
             previous_event = EventOverlap(EventHandle())
+        # Build configs; combine uses combine_config, backward dispatch uses dispatch_config.
+        combine_config = _build_deepep_config(
+            Buffer.get_combine_config(group.size()), fused_a2a_config
+        )
+        dispatch_config = _build_deepep_config(
+            Buffer.get_dispatch_config(group.size()), fused_a2a_config
+        )
         buffer = get_buffer(group, get_hidden_bytes(x))
         combined_x, _, after_event = buffer.combine(
             x,
             handle=handle,
+            config=combine_config,
             async_finish=async_finish,
             previous_event=previous_event,
             allocate_on_comm_stream=allocate_on_comm_stream,
@@ -187,6 +263,7 @@ def forward(ctx, x, group, handle, async_finish=False, allocate_on_comm_stream=F
         ctx.group = group
         ctx.async_finish = async_finish
         ctx.allocate_on_comm_stream = allocate_on_comm_stream
+        ctx.dispatch_config = dispatch_config
         return combined_x, None
 
     @staticmethod
@@ -199,14 +276,15 @@ def backward(ctx, grad_output, previous_event=None):
         grad_x, _, _, _, _, after_event = buffer.dispatch(
             grad_output.contiguous(),
             handle=ctx.handle,
+            config=ctx.dispatch_config,
             previous_event=previous_event,
             async_finish=ctx.async_finish,
             allocate_on_comm_stream=ctx.allocate_on_comm_stream,
         )
         # Make sure current stream is synchronized
         if ctx.async_finish:
             after_event.current_stream_wait()
-        return grad_x, None, None, None, None
+        return grad_x, None, None, None, None, None
 
 
 if HAVE_DEEP_EP:
@@ -219,6 +297,7 @@ def fused_dispatch(
         group,
         async_finish=False,
         allocate_on_comm_stream=False,
+        config=None,
     ):
         """Perform fused dispatch operation if deep_ep is available.
 
@@ -241,9 +320,11 @@ def fused_dispatch(
             group,
             async_finish,
             allocate_on_comm_stream,
+            config,
         )
 
-    def fused_combine(x, group, handle, async_finish=False, allocate_on_comm_stream=False):
+    def fused_combine(x, group, handle, async_finish=False, allocate_on_comm_stream=False,
+                      config=None):
         """Perform fused combine operation if deep_ep is available.
 
         Args:
@@ -255,7 +336,7 @@ def fused_combine(x, group, handle, async_finish=False, allocate_on_comm_stream=
         Returns:
             Result of FusedCombine
         """
-        return FusedCombine.apply(x, group, handle, async_finish, allocate_on_comm_stream)
+        return FusedCombine.apply(x, group, handle, async_finish, allocate_on_comm_stream, config)
 
     def set_deepep_num_sms(num_sms):
         """Sets the number of SMs to use for DeepEP"""

megatron/core/transformer/moe/fused_a2a_config.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""
+FusedA2AConfig dataclass for user-tunable fused all-to-all MoE parameters.
+"""
+from dataclasses import dataclass, field
+from typing import Optional
+
+
+@dataclass
+class FusedA2AConfig:
+    """
+    Configuration for fused all-to-all MoE (FusedDispatch/FusedCombine).
+
+    Fields:
+        chunk_size: Optional[int] - Chunk size for all-to-all (must be positive if set)
+        num_sms: Optional[int] - Number of SMs for kernel (must be positive and even if set)
+        # Future tunables can be added here
+
+    Precedence: CLI > ENV > CONFIG FILE > DEFAULTS
+    """
+    chunk_size: Optional[int] = None
+    num_sms: Optional[int] = None
+    # Future tunables can be added here
+
+    def validate(self):
+        if self.chunk_size is not None:
+            if not (self.chunk_size > 0):
+                raise ValueError(f"chunk_size must be positive, got {self.chunk_size}")
+        if self.num_sms is not None:
+            if not (self.num_sms > 0):
+                raise ValueError(f"num_sms must be positive, got {self.num_sms}")
+            # DeepEP's Buffer.set_num_sms asserts new_num_sms % 2 == 0
+            # and the C++ kernel asserts config.num_sms % 2 == 0. An odd value
+            # would crash deep inside the kernel with an opaque assertion;
+            # fail fast at validate_args time instead.
+            if self.num_sms % 2 != 0:
+                raise ValueError(
+                    f"num_sms must be even (DeepEP requirement), got {self.num_sms}"
+                )
+
+    @staticmethod
+    def from_dict(cfg: dict) -> 'FusedA2AConfig':
+        allowed = {'chunk_size', 'num_sms'}
+        unknown = set(cfg.keys()) - allowed
+        if unknown:
+            raise ValueError(f"Unknown FusedA2AConfig keys: {unknown}")
+        return FusedA2AConfig(
+            chunk_size=cfg.get('chunk_size'),
+            num_sms=cfg.get('num_sms'),
+        )
+
+    def __repr__(self):
+        return f"FusedA2AConfig(chunk_size={self.chunk_size}, num_sms={self.num_sms})"

megatron/core/transformer/moe/fused_a2a_config_loader.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+"""
+FusedA2A config loader and resolver for Megatron-LM MoE fused all-to-all.
+"""
+import os
+import json
+from typing import Optional
+
+try:
+    import yaml
+    HAVE_YAML = True
+except ImportError:
+    HAVE_YAML = False
+
+from .fused_a2a_config import FusedA2AConfig
+
+
+def load_a2a_config_from_file(path: str) -> dict:
+    """
+    Load config from JSON or YAML file.
+    Raises ValueError on error.
+    """
+    if path.endswith('.json'):
+        with open(path, 'r') as f:
+            return json.load(f)
+    elif path.endswith(('.yaml', '.yml')):
+        if not HAVE_YAML:
+            raise ImportError('pyyaml is required for YAML config files')
+        with open(path, 'r') as f:
+            return yaml.safe_load(f)
+    else:
+        raise ValueError(f"Unsupported config file extension: {path}")
+
+def resolve_fused_a2a_config_from_sources(cli_args=None, env=os.environ, config_file_path=None) -> FusedA2AConfig:
+    """
+    Resolve FusedA2AConfig from CLI args, environment, and config file.
+    Precedence: CLI > ENV > CONFIG FILE > DEFAULTS.
+    Raises ValueError on any invalid or unknown keys.
+    """
+    file_cfg = {}
+    if config_file_path:
+        file_cfg = load_a2a_config_from_file(config_file_path)
+    env_cfg = {}
+    if env.get('MOE_A2A_CHUNK_SIZE'):
+        env_cfg['chunk_size'] = int(env['MOE_A2A_CHUNK_SIZE'])
+    if env.get('MOE_A2A_NUM_SMS'):
+        env_cfg['num_sms'] = int(env['MOE_A2A_NUM_SMS'])
+    cli_cfg = {}
+    if cli_args is not None:
+        if getattr(cli_args, 'moe_a2a_chunk_size', None) is not None:
+            cli_cfg['chunk_size'] = cli_args.moe_a2a_chunk_size
+        if getattr(cli_args, 'moe_a2a_num_sms', None) is not None:
+            cli_cfg['num_sms'] = cli_args.moe_a2a_num_sms
+    merged = {**file_cfg, **env_cfg, **cli_cfg}
+    config = FusedA2AConfig.from_dict(merged)
+    config.validate()
+    return config