CuteDSL MegaMoe Optimization for DSV41

docs/design_docs/moe_ep_sm107_qualification.md

@@ -178,9 +178,12 @@ Per-call tensors on `MoEEpTensors` override the matching config tensor. An
 omitted override restores the config default, or one if the config also omits
 it. Scalars have shape `[E_local]`, are contiguous FP32 tensors on the input
 device, and follow local expert order. Their device assertions remain active
-in graph replay. Values are copied into stable workspace buffers on every
-forward, including empty batches, so sequential pooled layers do not share
-stale values. Update captured source tensor contents in place for graph replay.
+in graph replay. Contiguous FP32 per-call tensors are bound as launch
+arguments on every forward, including empty batches, so sequential pooled
+layers do not share stale values; keep them alive while launches or captured
+graphs use them. Other values are copied into stable workspace buffers, and an
+omitted scalar resets its buffer to one. Update captured source tensor
+contents in place for graph replay.
 
 For `quantize_input=False`, input bytes and scales are copied unchanged;
 `input_norm_const` does not requantize them. The caller must supply alphas

flashinfer/moe_ep/backends/mega/kernel/sm107/nvfp4_nvfp4_bf16_cutedsl/backend.py

@@ -231,12 +231,7 @@ def stage_inputs(
         validate_routing_values(
             t.topk_ids, t.topk_weights, workspace.config.num_total_experts
         )
-        for name, value in self._scalar_values(t).items():
-            target = getattr(workspace, name)
-            if value is None:
-                target.fill_(1.0)
-            else:
-                target.copy_(value)
+        self._stage_scalars(t, workspace)
         if quantize_input:
             staged = stage_mega_moe_inputs(
                 t.hidden_states,
@@ -268,6 +263,36 @@ def stage_inputs(
             workspace.topk_idx[num_tokens:capacity].fill_(-1)
         workspace.note_staged_tokens(num_tokens)
 
+    def _stage_scalars(self, t: "MoEEpTensors", workspace: Any) -> None:
+        """Per-call NVFP4 epilogue scalars.
+
+        Contiguous fp32 ``[E_local]`` tensors on the workspace device are BOUND as
+        the launch arguments (``workspace.bind_scalars``): no per-call DtoD copies
+        into the pooled workspace tensors (3 launches per MoE layer). Any other
+        value is copied into the workspace tensor and ``None`` resets it to ones,
+        as before; the launch then uses the workspace tensor."""
+        values = self._scalar_values(t)
+        bind = hasattr(workspace, "bind_scalars")
+        bound = {}
+        for name, value in values.items():
+            target = getattr(workspace, name)
+            if (
+                bind
+                and isinstance(value, torch.Tensor)
+                and value.dtype == torch.float32
+                and value.is_contiguous()
+                and value.device == target.device
+                and value.shape == target.shape
+            ):
+                bound[name] = value
+                continue
+            if value is None:
+                target.fill_(1.0)
+            else:
+                target.copy_(value)
+        if bind:
+            workspace.bind_scalars(bound)
+
     def compute(
         self,
         workspace: Any,

flashinfer/moe_ep/backends/mega/kernel/sm107/validation.py

@@ -20,6 +20,12 @@ def make_workspace_config(
     tuning = {name: getattr(k, name) for name in KNOB_KEYS if hasattr(k, name)}
     tuning["reduce_topk_in_kernel"] = k.in_kernel_fc2_reduce
     tuning.update(overrides)
+    if k.max_sm_count is not None:
+        # An SM budget (e.g. to leave SMs to a concurrent shared-expert MLP on a
+        # side stream) needs a uniform cluster launch: the mixed-CGA fallback
+        # occupancy is resolved from the hardware probe, so cached / heuristic
+        # knobs that name a fallback cluster are launched uniformly instead.
+        tuning["fallback_cluster_shape_mn"] = None
     for name in (
         "mma_tiler_mnk",
         "cluster_shape_mn",

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/VENDOR.md

@@ -163,14 +163,36 @@ are outside the reset prefix. GenPhase's reset bound depends on both top-k and
 local expert count; large-prefix constructor checks must stay within its
 1024-token capacity limit. EP16 construction is a layout test, not EP16 execution.
 
+One more FlashInfer-owned kernel fix, pending upstream (the pinned commit and
+upstream `main` up to `413da003` have the same code):
+
+- `src/sources/kernel_src/schedulers/fc12_scheduler.py`,
+  `PhaseInterleavedFc12Scheduler._wait_for_fc1_prologue_claims`: the
+  phase-interleave scheduler no longer sets `fc1_exhausted` from its CTA-local
+  peek of the global FC1 work-ID counter. Every CTA of a cluster runs the claim
+  state machine itself and interprets the leader's broadcast work ID against
+  its own stream choice, so that state must only change on cluster-uniform
+  inputs. When the runtime FC1 stream is shorter than the static FC1 watermark
+  (tiny per-rank work, e.g. DP dummy batches), the last terminal FC1 claims of
+  other clusters can push the counter across the watermark between the peeks
+  of two CTAs of one cluster (late-starting clusters make this likely). The
+  CTAs then disagreed on `fc1_exhausted`, mapped the next broadcast ID to
+  different phases, and the cluster deadlocked (CTAs leaving the claim loop
+  while the leader keeps claiming, or 2-CTA MMA / multicast pairs on different
+  tiles): one rank 2 CTAs short of the tail grid sync, its peers waiting in
+  the NVLink drain barrier. The peek now only shortens this CTA's wait; FC1
+  exhaustion is learned from a claimed ID like any other stream end, at the
+  cost of at most one extra terminal FC1 claim per cluster. Re-apply this fix
+  when the drop is re-exported until upstream carries it.
+
 FlashInfer-owned `shim/block_scaled.py` admits the geometries listed above and
 omits BF16 masked-slot preclear only when the kernel advertises masked reduction.
 Quantized combine and in-kernel reduction retain their existing clearing rules.
 This shim change is outside the exported payload. The local patch inventory
 above is separate from the historical upstream qualification evidence; run the
 registered strict single/multirank suites when updating or rebasing these patches.
 All remaining exported implementation files retain the pinned export, apart
-from the previously documented comment.
+from the previously documented comment and the scheduler fix above.
 
 ## Related trees
 

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/shim/autotune.py

@@ -272,7 +272,9 @@ def barrier():
         trial.topk_weights.copy_(symm_buffer.topk_weights)
         if cfg.quant_kind == "nvfp4":
             for name in ("fc1_alpha", "fc2_alpha", "fc1_norm_const"):
-                getattr(trial, name).copy_(getattr(symm_buffer, name))
+                # the scalars the workspace's launches use (bound per-call tensors by
+                # default, where the workspace tensors stay unwritten)
+                getattr(trial, name).copy_(symm_buffer._scalar_tensor(name))
         trial.note_staged_tokens(num_tokens)
         barrier()
         for _ in range(warmup_iters):

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/shim/block_scaled.py

@@ -52,6 +52,8 @@
 
 TransformedBlockScaledWeights = Tuple[torch.Tensor, torch.Tensor]
 
+_SCALAR_NAMES = ("fc1_alpha", "fc2_alpha", "fc1_norm_const")
+
 
 def _require_sm107() -> None:
     cc = torch.cuda.get_device_capability()
@@ -489,8 +491,9 @@ def __init__(self, config: Sm107BlockScaledMoeConfig) -> None:
             (tokens, cfg.num_topk), -1, dtype=torch.int32, device="cuda"
         )
         # Stable local-expert pointers allow per-call scaling updates in graphs.
-        # Defaults and overrides are copied on every backend staging call.
-        for name in ("fc1_alpha", "fc2_alpha", "fc1_norm_const"):
+        # Per-call tensors are bound as launch arguments (bind_scalars); other
+        # values and defaults are copied in by the backend staging call.
+        for name in _SCALAR_NAMES:
             setattr(
                 self,
                 name,
@@ -565,6 +568,9 @@ def __init__(self, config: Sm107BlockScaledMoeConfig) -> None:
         ] = None
         self._staged_tokens: Optional[int] = None
         self._destroyed = False
+        # Per-call NVFP4 epilogue scalars bound straight into the launch (no copy
+        # into the workspace tensors); a missing name = the workspace's own tensor.
+        self._bound_scalars: dict = {}
 
     @staticmethod
     def _build_kernel(cfg: Sm107BlockScaledMoeConfig):
@@ -678,6 +684,24 @@ def note_staged_tokens(self, num_tokens: int) -> None:
     def staged_tokens(self) -> Optional[int]:
         return self._staged_tokens
 
+    def bind_scalars(self, values: dict) -> None:
+        """Bind per-call NVFP4 scalar tensors (``fc1_alpha`` / ``fc2_alpha`` /
+        ``fc1_norm_const``: contiguous fp32 ``[experts_per_rank]`` CUDA tensors) as
+        the launch arguments, instead of copying them into the workspace tensors
+        every call. ``None`` (or a missing name) selects the workspace tensor.
+        The tensors must stay alive while launches (or captured graphs) use them."""
+        bound = {}
+        for name in _SCALAR_NAMES:
+            value = values.get(name)
+            if value is not None:
+                bound[name] = value
+        self._bound_scalars = bound
+
+    def _scalar_tensor(self, name: str):
+        bound = getattr(self, "_bound_scalars", None) or {}
+        value = bound.get(name)
+        return value if value is not None else getattr(self, name)
+
     def _peer_mapper(self):
         from sources.communication.nvlink_domain.symmetric_buffer import (
             SymmetricBufferHost,
@@ -711,8 +735,8 @@ def _runtime_kwargs(
             "peer_rank_ptr_mapper_host": self._peer_mapper(),
             "stream": _cu_stream(stream),
             **{
-                name: _to_cute(getattr(self, name), assumed_align=4)
-                for name in ("fc1_alpha", "fc2_alpha", "fc1_norm_const")
+                name: _to_cute(self._scalar_tensor(name), assumed_align=4)
+                for name in _SCALAR_NAMES
                 if self.config.quant_kind == "nvfp4"
             },
         }
@@ -739,6 +763,13 @@ def launch(
             transformed_l1[1].data_ptr(),
             transformed_l2[0].data_ptr(),
             transformed_l2[1].data_ptr(),
+            # Bound per-call scalars: a new tensor re-binds the launch arguments
+            # (host-only, no recompile).
+            *(
+                self._scalar_tensor(name).data_ptr()
+                for name in _SCALAR_NAMES
+                if self.config.quant_kind == "nvfp4"
+            ),
             torch.cuda.current_stream().cuda_stream,
         )
         if self._compiled is None or self._launch_key != key:
@@ -784,6 +815,7 @@ def destroy(self) -> None:
         self._launch_weights = None
         self._pre_reduced_activation = None
         self._pre_reduced_activation_sf = None
+        self._bound_scalars = {}
         for name in ("x", "x_sf", "topk_weights", "shared_workspace"):
             comm.free_sym_tensor(getattr(self, name, None))
             setattr(self, name, None)

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/shim/correctness.py

@@ -10,6 +10,16 @@
 from .kernel_helpers import compute_megamoe_reference_sm107_block_scaled
 
 
+def _launch_scalar(symm_buffer, name):
+    """The NVFP4 epilogue scalar tensor the next launch uses: the per-call tensor bound
+    with ``bind_scalars`` (the workspace tensor is then never written) or the
+    workspace's own tensor."""
+    get = getattr(symm_buffer, "_scalar_tensor", None)
+    if get is not None and getattr(symm_buffer, name, None) is not None:
+        return get(name)
+    return getattr(symm_buffer, name, None)
+
+
 def sampled_reference(
     symm_buffer, l1, l2, num_tokens: int, *, process_group=None, sample_size=64
 ):
@@ -81,9 +91,9 @@ def sampled_reference(
             return_fp32=True,
             situ_beta=cfg.situ_beta,
             situ_linear_beta=cfg.situ_linear_beta,
-            fc1_alpha=getattr(symm_buffer, "fc1_alpha", None),
-            fc2_alpha=getattr(symm_buffer, "fc2_alpha", None),
-            fc1_norm_const=getattr(symm_buffer, "fc1_norm_const", None),
+            fc1_alpha=_launch_scalar(symm_buffer, "fc1_alpha"),
+            fc2_alpha=_launch_scalar(symm_buffer, "fc2_alpha"),
+            fc1_norm_const=_launch_scalar(symm_buffer, "fc1_norm_const"),
         )
         if cfg.world_size > 1:
             dist.reduce(

flashinfer/moe_ep/kernel_src/sm107/next_cutedsl_megamoe/src/sources/kernel_src/schedulers/fc12_scheduler.py

@@ -803,17 +803,24 @@ def _wait_for_fc1_prologue_claims(
         work_id_worker: NonClcMixedCgaSchedulerWorker,
         task_mapping_state,
     ) -> Boolean:
-        """Wait until the static FC1 watermark or the smaller runtime stream end is claimed."""
+        """Wait until the static FC1 watermark or the smaller runtime stream end is claimed.
+
+        FlashInfer local fix (VENDOR.md "Local differences"): the peek below is one CTA-local load of a counter
+        that other clusters keep incrementing, so CTAs of one cluster (and the physical clusters of one fallback
+        group) can observe different values. The observation therefore only decides how long this CTA waits; it
+        never changes the scheduler state. Upstream set ``fc1_exhausted`` when the peek missed the watermark and the
+        runtime FC1 stream ends before it: a CTA that peeked just before the last terminal FC1 claims crossed the
+        watermark then disagreed with its cluster peers on ``fc1_exhausted``, the peers then mapped the same
+        broadcast work ID to different phases (each CTA uses its own stream choice), and the cluster deadlocked
+        (some CTAs left the claim loop while the leader kept claiming, or the 2-CTA MMA / multicast pairs ran
+        different tiles). FC1 exhaustion is now always learned from a claimed ID, which every CTA sees alike.
+        """
         if not fc1_exhausted:
             fc1_counter_pointer = work_id_worker.get_atomic_counter_pointer(Int32(0))
             minimum_claim_count = Int32(self.minimum_global_fc1_claims)
             if not spin_peek(fc1_counter_pointer, lambda value: value >= minimum_claim_count):
-                claim_target, stream_ends_before_target = resolve_phase_interleaved_fc1_claim_target(
-                    minimum_claim_count, task_mapping_state
-                )
+                claim_target, _ = resolve_phase_interleaved_fc1_claim_target(minimum_claim_count, task_mapping_state)
                 spin_wait(fc1_counter_pointer, lambda value: value >= claim_target)
-                if stream_ends_before_target:
-                    fc1_exhausted = Boolean(True)
         return fc1_exhausted
 
     @cute.jit

flashinfer/quantization/kernels/nvfp4_quantize.py

@@ -2507,6 +2507,40 @@ def _get_compiled_kernel_nvfp4_tma(
     return compiled_kernel, kernel_obj.rows_per_block
 
 
+def _nvfp4_outputs(
+    out: Optional[Tuple[torch.Tensor, torch.Tensor]],
+    m: int,
+    k: int,
+    scale_output_size: int,
+    device: torch.device,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Fresh ``(fp4 [m, k/2], flat scale [scale_output_size])`` buffers, or views of the
+    caller-owned ``out`` destinations (validated: dtype, contiguity, device, size)."""
+    if out is None:
+        return (
+            torch.empty(m, k // 2, dtype=torch.uint8, device=device),
+            torch.empty(scale_output_size, dtype=torch.uint8, device=device),
+        )
+    fp4_out, sf_out = out
+    for name, t in (("fp4_out", fp4_out), ("sf_out", sf_out)):
+        if t.dtype != torch.uint8 or not t.is_contiguous() or t.device != device:
+            raise ValueError(
+                f"nvfp4_quantize_cute_dsl out={name} must be a contiguous uint8 "
+                f"tensor on {device}"
+            )
+    if tuple(fp4_out.shape) != (m, k // 2):
+        raise ValueError(
+            f"nvfp4_quantize_cute_dsl out fp4 shape {tuple(fp4_out.shape)} != {(m, k // 2)}"
+        )
+    flat = sf_out.view(-1)
+    if flat.numel() < scale_output_size:
+        raise ValueError(
+            f"nvfp4_quantize_cute_dsl out scale buffer holds {flat.numel()} bytes, "
+            f"the selected kernel writes {scale_output_size}"
+        )
+    return fp4_out, flat[:scale_output_size]
+
+
 @flashinfer_api
 def nvfp4_quantize_cute_dsl(
     input: torch.Tensor,
@@ -2515,6 +2549,7 @@ def nvfp4_quantize_cute_dsl(
     enable_pdl: bool | None = None,
     # Appended last so existing positional construction keeps working.
     nvfp4_4over6: Optional[NVFP44Over6Config] = _UNSET,
+    out: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     r"""Quantize input tensor to NVFP4 format using the CuTe-DSL kernel.
 
@@ -2551,6 +2586,14 @@ def nvfp4_quantize_cute_dsl(
         - :class:`NVFP44Over6Config`: on with exactly that recipe. The
           environment is ignored.
 
+    out : (torch.Tensor, torch.Tensor), optional
+        Caller-owned destinations ``(fp4_out, sf_out)``: ``fp4_out`` a contiguous
+        ``uint8 [M, K/2]`` view, ``sf_out`` a contiguous ``uint8`` buffer whose
+        flat prefix receives the scale factors (it must hold at least the padded
+        scale size of the selected path; with the linear layout and M below the
+        TMA crossover that is exactly ``M * K/16``). The kernels write straight
+        into them (no extra copy); the returned tensors are views of them.
+
     Returns
     -------
     Tuple[torch.Tensor, torch.Tensor]
@@ -2598,6 +2641,12 @@ def nvfp4_quantize_cute_dsl(
             padded_sf_cols = num_sf_blocks_per_row
         else:
             padded_sf_cols = ((num_sf_blocks_per_row + 3) // 4) * 4
+        if out is not None:
+            # Same contract as a launch: validated views of the caller's buffers.
+            fp4_output, scale_output = _nvfp4_outputs(
+                out, m, k, m * padded_sf_cols, input.device
+            )
+            return fp4_output, scale_output.view(m, padded_sf_cols)
         return (
             torch.empty((m, k // 2), dtype=torch.uint8, device=input.device),
             torch.empty((m, padded_sf_cols), dtype=torch.uint8, device=input.device),
@@ -2679,9 +2728,8 @@ def nvfp4_quantize_cute_dsl(
             (padded_m + rows_per_block - 1) // rows_per_block, tma_target_grid
         )
 
-        fp4_output = torch.empty(m, k // 2, dtype=torch.uint8, device=input.device)
-        scale_output = torch.empty(
-            scale_output_size, dtype=torch.uint8, device=input.device
+        fp4_output, scale_output = _nvfp4_outputs(
+            out, m, k, scale_output_size, input.device
         )
 
         kernel_fn(
@@ -2740,9 +2788,8 @@ def nvfp4_quantize_cute_dsl(
             target_grid,
         )
 
-        fp4_output = torch.empty(m, k // 2, dtype=torch.uint8, device=input.device)
-        scale_output = torch.empty(
-            scale_output_size, dtype=torch.uint8, device=input.device
+        fp4_output, scale_output = _nvfp4_outputs(
+            out, m, k, scale_output_size, input.device
         )
 
         kernel_fn(
@@ -2763,9 +2810,8 @@ def nvfp4_quantize_cute_dsl(
         padded_sf_cols = ((num_sf_blocks_per_row + 3) // 4) * 4
         scale_output_size = padded_m * padded_sf_cols
 
-        fp4_output = torch.empty(m, k // 2, dtype=torch.uint8, device=input.device)
-        scale_output = torch.empty(
-            scale_output_size, dtype=torch.uint8, device=input.device
+        fp4_output, scale_output = _nvfp4_outputs(
+            out, m, k, scale_output_size, input.device
         )
 
         use_tile = (