fix(moe_ep): ignore non-finite scores on masked Rubin routes

CLAUDE.md

@@ -788,7 +788,8 @@ Used by `flashinfer.trace` / `fi_trace`.
 | `FLASHINFER_VALIDATE_INPUTS` | `0` | `flashinfer/mla/_core.py` (MLA wrapper) | Non-zero / non-empty value enables defensive input validation inside the MLA wrapper. Device-synchronizing checks are skipped during CUDA Graph capture. Adds host-side overhead; intended for debugging. |
 | `FLASHINFER_AUTOTUNER_LOAD_FROM_FILE` | `0` | `flashinfer/autotuner/autotuner.py` | `1` loads previously serialized autotune results from disk instead of re-running the search. |
 | `FLASHINFER_DIST_AWARE_AUTOTUNE` | `0` | `flashinfer/fused_moe/da_config.py` | `1` enables experimental distribution-aware autotune and kernel dispatch (TRT-LLM MoE only). |
-| `FLASHINFER_SM107_FUSED_ROUTE_CHECK` | `1` | `flashinfer/moe_ep/backends/mega/kernel/sm107/validation.py` | `0` uses Torch routing checks instead of the fused checker. Both paths check expert IDs, duplicates, and finite scores, including during graph replay. |
+| `FLASHINFER_SM107_MEGA_CHECK_ROUTES` | `1` | `flashinfer/moe_ep/backends/mega/kernel/sm107/validation.py` | `0` skips routing value checks for caller-validated inputs. Metadata checks and masked-score zeroing remain active. Set before warmup and graph capture; changing it requires recapturing graphs. |
+| `FLASHINFER_SM107_FUSED_ROUTE_CHECK` | `1` | `flashinfer/moe_ep/backends/mega/kernel/sm107/validation.py` | `0` uses Torch routing checks instead of the fused checker. When routing checks are enabled, both paths check expert IDs, duplicates, and finite scores, including during graph replay. |
 | `FLASHINFER_SM107_PREP_TVM_FFI` | `1` | `flashinfer/moe_ep/backends/mega/kernel/sm107/routing_validation.py` | `0` launches the fused routing checker through CuTe's DLPack bindings instead of TVM-FFI. Set before warmup; each mode needs its own compiled kernel. |
 | `FLASHINFER_MEGAMOE_LOG_COMPILE_SPEC` | unset | `flashinfer/moe_ep/kernel_src/sm120/split_cutedsl_megakernel/shim/runtime.py` | Exactly `1` prints MegaMoE split-kernel compile choices, workspace memory diagnostics, and frontend compile/capture timings; other values leave these diagnostics off. |
 | `FLASHINFER_DA_DISTRIBUTIONS` | built-in distribution catalog | `flashinfer/fused_moe/da_config.py` | Comma-separated training distributions used by the experimental TRT-LLM distribution-aware MoE autotuner. |

docs/design_docs/moe_ep_sm107_qualification.md

@@ -65,11 +65,28 @@ topology. See the [NVSHMEM environment variable reference](https://docs.nvidia.c
 | Prequantized inputs | Matching data format; scales are `[T, H/16]` E4M3 for NVFP4 or `[T, H/32]` E8M0 for MXFP8; uint8 scale storage is interpreted as raw bytes. Pass the logical scale columns, excluding internal workspace communication padding. |
 | Activation | SwiGLU by default; SiTU requires both positive, finite beta parameters and excludes gate/up clamps |
 | Normalization | NVFP4 accepts positive, finite local-expert FP32 vectors for `fc1_alpha`, `fc2_alpha`, and `fc1_norm_const`. Per-call values override config defaults; if both omit a value, staging uses one. MXFP8 and MXFP4/MXFP8 do not accept these scalars. |
-| Routing values | Unique valid expert IDs per token or `-1` masked slots; finite scores; repeated masked slots are allowed |
+| Routing values | Unique valid expert IDs per token or `-1` masked slots; finite scores for valid routes; repeated masked slots are allowed |
 | Output | BF16; owned tensor by default; workspace views expire on the next workspace use or destruction |
 | Graphs and pooling | Warm up eagerly on every rank before capture; sequential, stream-ordered use of a shared workspace |
 | Determinism | In-kernel FC2 reduction is opt-in and requires early routing weights; repeatability must be characterized separately |
 
+Scores for masked (`-1`) routes may be NaN or infinite. Staging writes zero
+for those scores without modifying the caller's routing tensors.
+
+Prequantized inputs use a route-staging kernel to copy IDs, convert scores to
+FP32 and zero masked scores. Activation and scale bytes are copied unchanged.
+Eligible BF16 inputs use fused quantization and staging for NVFP4 or MXFP8;
+unsupported layouts and score dtypes retain the Torch path. Set
+`FLASHINFER_MEGA_FUSED_STAGE=0` to use Torch staging. Routing validation remains
+separate and runs before either staging path by default.
+
+Set `FLASHINFER_SM107_MEGA_CHECK_ROUTES=0` to skip routing value checks when
+the caller guarantees valid routing. This skips ID-range, duplicate-ID and
+finite-score checks; it does not sanitize non-finite scores on active routes.
+Metadata checks, staging and masked-score zeroing remain active. Set the
+variable before warmup and graph capture; existing graphs must be recaptured
+to change their validation behavior.
+
 Routing value checks use device assertions so they remain active on graph
 replay without a host synchronization. Invalid routing can invalidate the
 worker's CUDA context, as with other invalid CUDA indexing operations.
@@ -90,8 +107,8 @@ Change tensor contents in place; recapture for a new shape/count, or mask
 inactive rows while retaining the captured shape.
 
 BF16 Rubin kernels, training, local fused-routing MegaMoE, FP4 activations
-with E8M0 scales, and cross-node communication are not part of this implementation. BF16 input
-staging uses Torch operations; its cost is included in full-forward measurements.
+with E8M0 scales, and cross-node communication are not part of this implementation.
+Input staging is included in full-forward measurements.
 
 ## SiTU and prequantized weights
 

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

@@ -253,8 +253,15 @@ def stage_inputs(
             workspace.x_sf[:num_tokens, :sf_cols].view(torch.uint8).copy_(
                 t.scales.view(torch.uint8)
             )
-            workspace.topk_idx[:num_tokens].copy_(t.topk_ids.to(torch.int32))
-            workspace.topk_weights[:num_tokens].copy_(t.topk_weights.to(torch.float32))
+            from ..routing_staging import stage_routes
+
+            stage_routes(
+                t.topk_ids,
+                t.topk_weights,
+                workspace.topk_idx,
+                workspace.topk_weights,
+                workspace=workspace,
+            )
         if num_tokens < capacity:
             workspace.topk_idx[num_tokens:capacity].fill_(-1)
         workspace.note_staged_tokens(num_tokens)

flashinfer/moe_ep/backends/mega/kernel/sm107/mxfp8_mxfp8_bf16_cutedsl/staging.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import os
+
 import torch
 
 
@@ -26,8 +28,7 @@ def stage_mega_moe_inputs(
 ) -> int:
     """bf16 ``hidden_states`` -> MXFP8 activation + E8M0 block scales.
 
-    Torch-composed staging (the ``next/`` drop has no fused staging kernel
-    yet). Returns the staged token count.
+    Returns the staged token count. Unsupported layouts use Torch staging.
     """
     from ......kernel_src.sm107.next_cutedsl_megamoe import (
         Mxfp8BlockSize,
@@ -44,6 +45,40 @@ def stage_mega_moe_inputs(
     if topk_weights.shape != topk_ids.shape:
         raise ValueError("topk_weights and topk_ids must have the same shape.")
 
+    if (
+        os.environ.get("FLASHINFER_MEGA_FUSED_STAGE", "1") != "0"
+        and hidden_states.is_cuda
+        and hidden_states.is_contiguous()
+        and hidden_states.data_ptr() % 16 == 0
+        and hidden % 128 == 0
+        and x_sf.shape[1] == hidden // 32
+        and topk_ids.shape[1] <= 128
+        and topk_ids.is_contiguous()
+        and topk_weights.is_contiguous()
+        and topk_ids.data_ptr() % 16 == 0
+        and topk_weights.data_ptr() % 16 == 0
+        and topk_weights.dtype in (torch.float16, torch.bfloat16, torch.float32)
+    ):
+        from ......kernel_src.sm100.cutedsl_megamoe import (
+            fused_quant_stage,
+            note_staged_tokens,
+        )
+
+        # Pooled workspaces may have been written by another staging path.
+        note_staged_tokens(topk_idx_out, capacity)
+        fused_quant_stage(
+            hidden_states,
+            topk_ids,
+            topk_weights,
+            x_fp8,
+            x_sf,
+            topk_idx_out,
+            topk_weights_out,
+            quant_type="mxfp8_e4m3" if kind == "mxfp4_mxfp8" else kind,
+            mxfp8_normal_scales=True,
+        )
+        return num_tokens
+
     q, sf = quantize_mxfp8_block32(hidden_states.to(torch.float32), _data_dtype(kind))
 
     hidden_sf_cols = ceil_div(hidden, Mxfp8BlockSize)
@@ -57,7 +92,10 @@ def stage_mega_moe_inputs(
     x_sf[:num_tokens].view(torch.uint8).zero_()
     x_sf[:num_tokens, :hidden_sf_cols].view(torch.uint8).copy_(sf.view(torch.uint8))
     topk_idx_out[:num_tokens].copy_(topk_ids.to(torch.int32))
-    topk_weights_out[:num_tokens].copy_(topk_weights.to(torch.float32))
+    # Masked scores may be non-finite; zero them before late-weight combine.
+    topk_weights_out[:num_tokens].copy_(
+        topk_weights.to(torch.float32).masked_fill(topk_ids == -1, 0)
+    )
     if num_tokens < capacity:
         topk_idx_out[num_tokens:capacity].fill_(-1)
     return num_tokens

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

@@ -260,8 +260,15 @@ def stage_inputs(
             workspace.x_sf[:num_tokens, :sf_cols].view(torch.uint8).copy_(
                 t.scales.view(torch.uint8)
             )
-            workspace.topk_idx[:num_tokens].copy_(t.topk_ids.to(torch.int32))
-            workspace.topk_weights[:num_tokens].copy_(t.topk_weights.to(torch.float32))
+            from ..routing_staging import stage_routes
+
+            stage_routes(
+                t.topk_ids,
+                t.topk_weights,
+                workspace.topk_idx,
+                workspace.topk_weights,
+                workspace=workspace,
+            )
         if num_tokens < capacity:
             workspace.topk_idx[num_tokens:capacity].fill_(-1)
         workspace.note_staged_tokens(num_tokens)

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

@@ -91,7 +91,10 @@ def stage_mega_moe_inputs(
     x_sf[:num_tokens].view(torch.uint8).zero_()
     x_sf[:num_tokens, :hidden_sf_cols].view(torch.uint8).copy_(sf.view(torch.uint8))
     topk_idx_out[:num_tokens].copy_(topk_ids.to(torch.int32))
-    topk_weights_out[:num_tokens].copy_(topk_weights.to(torch.float32))
+    # Masked scores may be non-finite; zero them before late-weight combine.
+    topk_weights_out[:num_tokens].copy_(
+        topk_weights.to(torch.float32).masked_fill(topk_ids == -1, 0)
+    )
     if num_tokens < capacity:
         topk_idx_out[num_tokens:capacity].fill_(-1)
     return num_tokens

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

@@ -0,0 +1,125 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+"""Copy validated routes into the MegaMoE workspace without mask temporaries."""
+
+from dataclasses import dataclass
+from functools import cache
+import os
+
+import torch
+
+
+@cache
+def _kernel_class():
+    import cuda.bindings.driver as cuda
+    import cutlass
+    import cutlass.cute as cute
+
+    class RouteStager:
+        def __init__(self, topk: int):
+            self.topk = topk
+
+        @cute.jit
+        def __call__(
+            self,
+            ids: cute.Tensor,
+            scores: cute.Tensor,
+            ids_out: cute.Tensor,
+            scores_out: cute.Tensor,
+            stream: cuda.CUstream,
+        ):
+            self.kernel(ids, scores, ids_out, scores_out).launch(
+                grid=((ids.shape[0] * self.topk + 255) // 256, 1, 1),
+                block=(256, 1, 1),
+                stream=stream,
+            )
+
+        @cute.kernel
+        def kernel(
+            self,
+            ids: cute.Tensor,
+            scores: cute.Tensor,
+            ids_out: cute.Tensor,
+            scores_out: cute.Tensor,
+        ):
+            tid, _, _ = cute.arch.thread_idx()
+            block, _, _ = cute.arch.block_idx()
+            index = block * 256 + tid
+            if index < ids.shape[0] * self.topk:
+                row, col = index // self.topk, index % self.topk
+                expert = cutlass.Int32(ids[row, col])
+                score = cutlass.Float32(0.0)
+                if expert != -1:
+                    score = cutlass.Float32(scores[row, col])
+                ids_out[row, col] = expert
+                scores_out[row, col] = score
+
+    return RouteStager
+
+
+@dataclass
+class _Entry:
+    key: tuple
+    inputs: tuple
+    arguments: tuple
+    compiled: object
+
+
+_COMPILED: dict[tuple, object] = {}
+
+
+def stage_routes(ids, scores, ids_out, scores_out, *, workspace) -> None:
+    """Stage live rows; the caller validates routes and clears the capacity tail."""
+    if (
+        not ids.is_cuda
+        or os.environ.get("FLASHINFER_MEGA_FUSED_STAGE", "1") == "0"
+        or scores.dtype not in (torch.float16, torch.bfloat16, torch.float32)
+    ):
+        live = ids.shape[0]
+        ids_out[:live].copy_(ids.to(torch.int32))
+        scores_out[:live].copy_(scores.float().masked_fill(ids == -1, 0))
+        return
+    if ids.numel() == 0:
+        return
+
+    import cuda.bindings.driver as cuda
+    import cutlass.cute as cute
+    import cutlass.torch as cutlass_torch
+
+    tensors = (ids, scores, ids_out, scores_out)
+    ffi = os.environ.get("FLASHINFER_SM107_PREP_TVM_FFI", "1") != "0"
+    key = (
+        ids.shape[1],
+        tuple((t.device, t.dtype, tuple(t.stride())) for t in tensors),
+        ffi,
+    )
+    entry = getattr(workspace, "_route_staging", None)
+    if (
+        entry is None
+        or entry.key != key
+        or any(a is not b for a, b in zip(entry.inputs, tensors, strict=True))
+    ):
+        compiled = _COMPILED.get(key)
+        arguments = tensors
+        if compiled is None or not ffi:
+            arguments = tuple(
+                cutlass_torch.from_dlpack(t, assumed_align=1).mark_layout_dynamic(
+                    leading_dim=1 if t.stride(1) == 1 else None
+                )
+                for t in tensors
+            )
+        if compiled is None:
+            if torch.cuda.is_current_stream_capturing():
+                raise RuntimeError("Warm route staging before CUDA graph capture")
+            compiled = cute.compile(
+                _kernel_class()(ids.shape[1]),
+                *arguments,
+                cuda.CUstream(torch.cuda.current_stream().cuda_stream),
+                options="--enable-tvm-ffi" if ffi else "",
+            )
+            _COMPILED[key] = compiled
+        entry = _Entry(key, tensors, tensors if ffi else arguments, compiled)
+        workspace._route_staging = entry
+    entry.compiled(
+        *entry.arguments, cuda.CUstream(torch.cuda.current_stream().cuda_stream)
+    )

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

@@ -51,12 +51,8 @@ def kernel(
             expert = cutlass.Int64(ids[row, col])
             score = cutlass.Float32(weights[row, col])
             valid = valid & (expert >= -1) & (expert < experts)
-            valid = (
-                valid
-                & (score == score)
-                & (score < float("inf"))
-                & (score > -float("inf"))
-            )
+            finite = (score == score) & (score < float("inf")) & (score > -float("inf"))
+            valid = valid & ((expert == -1) | finite)
             for other in cutlass.range_constexpr(self.topk):
                 if other < col:
                     duplicate = (expert >= 0) & (
@@ -154,5 +150,5 @@ def check_routing(
     )
     torch._assert_async(
         entry.result.all() if entry.result.numel() > 1 else entry.result,
-        "SM107 routes require unique expert IDs in [0, E), or -1, and finite scores",
+        "SM107 routes require unique expert IDs in [0, E), or -1, and finite scores for valid routes",
     )

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

@@ -179,13 +179,15 @@ def validate_forward_metadata(
 def validate_routing_values(
     topk_ids, topk_weights, num_experts: int, *, workspace=None
 ) -> None:
-    """Device assertions run before dispatch, including during graph replay.
+    """Device assertions run before dispatch unless explicitly disabled.
 
     Duplicate valid routes violate the router's workspace capacity proof.
-    -1 is a masked slot; all other IDs must be in range. A failed CUDA
-    assertion invalidates that worker's CUDA context, as with invalid indices
-    in other CUDA indexing operations.
+    -1 is a masked slot whose score is ignored; all other IDs must be in range
+    and have finite scores. A failed CUDA assertion invalidates that worker's
+    CUDA context, as with invalid indices in other CUDA indexing operations.
     """
+    if os.environ.get("FLASHINFER_SM107_MEGA_CHECK_ROUTES", "1") == "0":
+        return
     if (
         topk_ids.is_cuda
         and topk_ids.is_contiguous()
@@ -201,10 +203,11 @@ def validate_routing_values(
     ordered = topk_ids.sort(dim=1).values
     unique = (ordered[:, 1:] != ordered[:, :-1]) | (ordered[:, 1:] == -1)
     valid = ((topk_ids >= -1) & (topk_ids < num_experts)).all()
-    valid = valid & unique.all() & torch.isfinite(topk_weights).all()
+    finite = ((topk_ids == -1) | torch.isfinite(topk_weights)).all()
+    valid = valid & unique.all() & finite
     torch._assert_async(
         valid,
-        "SM107 routes require unique expert IDs in [0, E), or -1, and finite scores",
+        "SM107 routes require unique expert IDs in [0, E), or -1, and finite scores for valid routes",
     )