perf(moe_ep): add a separate scale ring to the SM90 NVFP4 CuTe DSL backend

docs/design_docs/moe_ep_sm90_nvfp4.md

@@ -178,3 +178,28 @@ five attempts retain the first block passing the rank-skew rule. Three rounds
 in one process are not three independent fresh-process runs. Use separate
 `torchrun` processes to measure process-to-process variation, and account for
 actual GPU clock and power states when interpreting results.
+
+## Two-slot scale ring
+
+`nvfp4_sf_stages=0` follows the AB pipeline's scale-buffer stage count.
+Eligible non-pingpong M256N64K256 tactics can instead use
+`nvfp4_sf_stages=2`, `compact_pull_buffer=True`, lane LUT mapping, and one
+or two LUT replicas. Two scale slots reduce shared-memory use; the supplied
+profiles use four AB stages. The AB and scale-ring barriers use their own
+slot counts. This changes buffering, not the numerical operation: same-core
+sf0/sf2 output remains byte-identical.
+
+The candidate table includes five measured SF2 profiles: R1/R2 variants
+of the T512 and T1024 seeds, and one R1 variant of the T2048 seed. Their
+admission uses geometry, protocol and shared-memory legality, without GPU
+product-name or live-token restrictions. T512/T1024 reference capacity
+buckets use R1 SF2 defaults; the other eleven reference defaults use SF0.
+Autotuning can select any legal profile by measurement, including with
+FC1 overlap explicitly disabled. Explicit token-back selection without
+`knobs` uses manual configuration rather than inheriting an SF2 default.
+
+A complete manual tactic can set `"nvfp4_sf_stages": 2`. The AB stage count
+must be a multiple of twice the scale-ring depth, and rounding must not
+reduce it below the corresponding SF0 stage count. Illegal resource
+combinations are rejected before launch. SF1 remains available for legal
+manual configurations that meet this rule.

flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/SKILL.md

@@ -137,6 +137,12 @@ not picked them up):
   builder shares its byte/row parameters between formats. Reapply these hooks
   when updating the drop, and
   compare complete MX and NV cubins before accepting a refactor.
+  The NV scale ring overrides `_weight_sf_stage_count`; reapply the base
+  hook for scale barrier storage and pipeline sizing independently of AB.
+  Shared MX producer/FC2 loops use `_weight_sf_decoupled` (default false)
+  and `_weight_sf_state` (default identity); the NV shim supplies only the
+  delayed scale-producer loop and early-release policy. Preserve these hooks
+  when reapplying the scale-ring extension.
 
 ## Layout
 

flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/hopper_nvfp4.py

@@ -115,8 +115,9 @@ def _validate_complete_nvfp4_tactic(
     """Validate one complete external fused tactic through the tuner schema.
 
     The MXFP4 schema fields are followed by the NVFP4 conversion fields
-    (``compact_pull_buffer``, ``nvfp4_lut_replicas``, ``nvfp4_lut_map``); a tactic without them selects the full pull buffer and the
-    8-replica lane-quad table.
+    (``compact_pull_buffer``, ``nvfp4_lut_replicas``, ``nvfp4_lut_map``,
+    ``nvfp4_sf_stages``); a tactic without them selects the full pull buffer,
+    8-replica lane-quad table and lockstep scale slots.
     """
 
     if not isinstance(knobs, dict):
@@ -147,6 +148,8 @@ class MegaMoEHopperNvfp4Config(MegaMoEHopperMxfp4Config):
     # Conversion-table replicas and lane map (tactic fields; values identical).
     nvfp4_lut_replicas: int = field(default=8, kw_only=True)
     nvfp4_lut_map: Literal["quad", "pair", "lane"] = field(default="quad", kw_only=True)
+    # Zero preserves lockstep weight scales; one or two use a separate ring.
+    nvfp4_sf_stages: int = field(default=0, kw_only=True)
     _weight_format = "NVFP4"
 
     def _validate_conversion_layout(self) -> None:
@@ -160,6 +163,8 @@ def _validate_active_dispatch_warps(self) -> None:
             raise ValueError("active_dispatch_warps must be 1, 2, or 4.")
 
     def _validate_weight_schedule(self) -> None:
+        if self.nvfp4_sf_stages and (self.mma_tiler_mnk[2] != 256 or self.pingpong):
+            raise ValueError("nvfp4_sf_stages requires MMA tile K=256 without pingpong")
         validate_nvfp4_fc1_overlap_override(
             self.nvfp4_fc1_k256_overlap,
             mma_tiler_mnk=self.mma_tiler_mnk,
@@ -173,6 +178,7 @@ def nvfp4_conversion(self) -> Nvfp4ConversionLayout:
             kperm=self.nvfp4_kperm,
             lut_replicas=self.nvfp4_lut_replicas,
             lut_map=self.nvfp4_lut_map,
+            sf_stages=self.nvfp4_sf_stages,
         )
 
 
@@ -202,6 +208,7 @@ def _conversion_tactic(self, kernel) -> Dict[str, Any]:
             "compact_pull_buffer": bool(kernel.token_comm.compact_pull_buffer),
             "nvfp4_lut_replicas": kernel.nvfp4_conversion.lut_replicas,
             "nvfp4_lut_map": kernel.nvfp4_conversion.lut_map,
+            "nvfp4_sf_stages": kernel.nvfp4_conversion.sf_stages,
         }
 
     _mega_key: Optional[tuple]

flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/nvfp4_candidate_data.py

@@ -123,27 +123,48 @@ def _tactic(**overrides):
     kperm=64,
 )
 
-# Capacity: (seed index, FC1 overlap default).
+# Capacity: (seed index, independent-scale-ring LUT replicas or None, overlap).
 _DEFAULTS = {
-    8: (0, False),
-    16: (1, True),
-    32: (2, True),
-    64: (3, False),
-    128: (4, True),
-    256: (5, False),
-    512: (6, True),
-    1024: (7, True),
-    2048: (8, True),
-    4096: (9, True),
-    8192: (10, True),
-    16384: (9, True),
-    32768: (11, True),
+    8: (0, None, False),
+    16: (1, None, True),
+    32: (2, None, True),
+    64: (3, None, False),
+    128: (4, None, True),
+    256: (5, None, False),
+    512: (6, 1, True),
+    1024: (7, 1, True),
+    2048: (8, None, True),
+    4096: (9, None, True),
+    8192: (10, None, True),
+    16384: (9, None, True),
+    32768: (11, None, True),
 }
 
+
+def _scale_ring_tactic(seed, replicas):
+    return dict(
+        seed,
+        compact_pull_buffer=True,
+        nvfp4_lut_replicas=replicas,
+        nvfp4_lut_map="lane",
+        nvfp4_sf_stages=2,
+    )
+
+
 NVFP4_DEFAULT_TACTICS_BY_TOKENS = {
     tokens: dict(
-        requested_tactic=dict(NVFP4_SF0_CANDIDATE_SEEDS[seed]),
+        requested_tactic=_scale_ring_tactic(NVFP4_SF0_CANDIDATE_SEEDS[seed], replicas)
+        if replicas is not None
+        else dict(NVFP4_SF0_CANDIDATE_SEEDS[seed]),
         requested_config=dict(kperm=64, nvfp4_fc1_k256_overlap=overlap),
     )
-    for tokens, (seed, overlap) in _DEFAULTS.items()
+    for tokens, (seed, replicas, overlap) in _DEFAULTS.items()
 }
+
+# Measured scale-ring choices. Admission depends on geometry, not GPU name,
+# live-token count, clamp, or overlap; construction checks actual resources.
+NVFP4_SF2_CANDIDATE_SEEDS = tuple(
+    _scale_ring_tactic(NVFP4_SF0_CANDIDATE_SEEDS[seed], replicas)
+    for seed, choices in ((6, (1, 2)), (7, (1, 2)), (9, (1,)))
+    for replicas in choices
+)

flashinfer/moe_ep/kernel_src/sm90/pull_style_cutedsl_megakernel/shim/nvfp4_conversion.py

@@ -3,7 +3,7 @@
 """SM90 NVFP4 x FP8 layouts and conversion primitives.
 
 Weights remain packed E2M1, with one raw nonnegative E4M3 scale per K16.
-The caller appends a replicated shared conversion table (8, 16 or 32 KiB, see
+The caller appends a replicated shared conversion table (1 to 32 KiB, see
 ``Nvfp4ConversionLayout``) to the scale ring and initializes it once before
 warp specialization. Promotion computes RN_E4M3(q*s/8) and supplies register
 operand A to K32 RS WGMMA; the existing epilogue applies 8G. MXFP4-prefixed
@@ -209,9 +209,13 @@ def nvfp4_initialize_lut(lookup: cute.Tensor, tidx: Int32):
 def nvfp4_word_smem(
     src: Int32, scale: Int32, upper: cutlass.Constexpr, lookup: cute.Tensor
 ):
-    # Lane-quad replica: valid for every table size (all replicas are equal).
+    # Lane-quad replica (all replicas are equal); tables with fewer than eight
+    # replicas fold the eight quads onto the ones they have.
     replicas = cute.size(lookup) // 128
-    bits = lookup[scale * Int32(replicas) + cute.arch.lane_idx() // Int32(4)]
+    replica = cute.arch.lane_idx() // Int32(4)
+    if cutlass.const_expr(replicas < 8):
+        replica = replica & Int32(replicas - 1)
+    bits = lookup[scale * Int32(replicas) + replica]
     lo = Int32(bits & cutlass.Uint64(0xFFFFFFFF))
     hi = Int32(bits >> cutlass.Uint64(32))
     selector = src & Int32(0x77777777)