feat(EPv2): forward a per-token scale row with dispatch

.github/workflows/ci_cco.yml

@@ -203,6 +203,45 @@ jobs:
             case \"\$out\" in *FAIL*) exit 1;; esac; \
             case \"\$out\" in *PASS*) ;; *) echo 'no PASS line: nothing ran'; exit 1;; esac"
 
+      - name: Run ops-v2 HIP backend with per-token scales (EP8, 8 ranks)
+        run: |
+          # SCALE_DIM>0 turns on the forwarded per-token scale row and makes
+          # test_op.py verify it through the routing reverse map (OP-SCALES). It is
+          # a separate step rather than an env added to the one above so the
+          # no-scale path keeps its own coverage.
+          $CT exec -e MORI_V2_KERNEL_BACKEND=hip -e SCALE_DIM=32 $CONTAINER bash -c "\
+            cd $GITHUB_WORKSPACE/tests/python/ops/dispatch_combine_v2 && \
+            out=\$(timeout 900 torchrun --standalone --nproc_per_node=8 test_op.py 2>&1) || \
+              { echo \"\$out\"; exit 1; }; \
+            echo \"\$out\"; \
+            case \"\$out\" in *FAIL*) exit 1;; esac; \
+            case \"\$out\" in *OP-SCALES*) ;; *) echo 'no OP-SCALES line: scales never ran'; exit 1;; esac"
+
+      - name: Run ops-v2 scales on a quantizing wire (fp8 and fp4 dispatch, EP8)
+        run: |
+          # The step above can only test scales on a BF16 wire: test_op.py gives
+          # dispatch and combine one dtype, and the HIP backend refuses an fp8
+          # combine (it sums). fp8/fp4 dispatch with a bf16 combine is what a real
+          # quantizing wire is, and it is the only configuration where the scale
+          # row is not decoration -- so it gets its own step.
+          #
+          # SCALE_DIM=224 on purpose: hidden/32 for hidden 7168, and NOT a multiple
+          # of 128, so it exercises the padding mori adds to the destination stride
+          # (EpScaleStride). A row that is already aligned would not.
+          #
+          # This runner has no gfx1250, so what actually runs here is the portable
+          # body -- the one that cannot be exercised by hand on the TDM boxes.
+          for disp in fp8 fp4; do
+            $CT exec -e MORI_V2_KERNEL_BACKEND=hip -e SCALE_DIM=224 -e HIDDEN=7168 \
+              -e DISP=$disp $CONTAINER bash -c "\
+              cd $GITHUB_WORKSPACE/tests/python/ops/dispatch_combine_v2 && \
+              out=\$(timeout 900 torchrun --standalone --nproc_per_node=8 test_asym_dtype.py 2>&1) || \
+                { echo \"\$out\"; exit 1; }; \
+              echo \"\$out\"; \
+              case \"\$out\" in *FAIL*) exit 1;; esac; \
+              case \"\$out\" in *ASYM-SCALES*) ;; *) echo 'no ASYM-SCALES line: scales never ran'; exit 1;; esac" || exit 1
+          done
+
       - name: Run ops-v2 backend parity (FlyDSL vs HIP, EP8)
         run: |
           $CT exec $CONTAINER bash -c "\

include/mori/ops/dispatch_combine_v2/ep_cfg.hpp

@@ -118,6 +118,9 @@ struct EpArgs {
   unsigned long long offDispOut = 0;    // T[maxRecv*hidden]     dispatch landing zone
   unsigned long long offOutTok = 0;     // T[maxRecv*hidden]     combine staging
   unsigned long long offXdb = 0;        // uint64[worldSize]     barrier slots
+  // uint8[maxRecv*scaleBytes]. Only read when Cfg.scaleBytes > 0; the arena does
+  // not carry the region otherwise, so this stays 0 and nothing dereferences it.
+  unsigned long long offOutScales = 0;
 
   // Which LSA rank this is. Runtime for the same reason: as a Cfg field it made
   // all eight ranks compile their own copy of an identical kernel.
@@ -126,6 +129,7 @@ struct EpArgs {
   const int* tokenIndices = nullptr;  // [numTokens * topk] expert ids, <0 drops
   const void* inpTokenBuf = nullptr;  // dispatch: source tokens; combine: post-expert tokens
   const float* weightsBuf = nullptr;  // [numTokens * topk]
+  const void* scalesBuf = nullptr;    // [numTokens * scaleBytes] per-token scale rows
   void* outTokenBuf = nullptr;        // combine output, local
   float* outWeightsBuf = nullptr;     // combine weight output, local
 
@@ -142,7 +146,7 @@ struct EpArgs {
 
 // The wire schema, generated from the field list rather than kept parallel to it.
 // The binding builds its ctypes struct from `name:tag` in this order and checks
-// sizeof -- which cannot see two same-type fields swapped, and 8 of the 22 are
+// sizeof -- which cannot see two same-type fields swapped, and 9 of the 24 are
 // bare pointers. So the static_asserts below take the offsets in SCHEMA order:
 // any disagreement with the declaration order stops the sequence increasing.
 #define MORI_EP_ARGS_FIELDS(X) \
@@ -155,10 +159,12 @@ struct EpArgs {
   X(offDispOut, "u64")         \
   X(offOutTok, "u64")          \
   X(offXdb, "u64")             \
+  X(offOutScales, "u64")       \
   X(rank, "i32")               \
   X(tokenIndices, "p")         \
   X(inpTokenBuf, "p")          \
   X(weightsBuf, "p")           \
+  X(scalesBuf, "p")            \
   X(outTokenBuf, "p")          \
   X(outWeightsBuf, "p")        \
   X(dispDestTokIdMap, "p")     \
@@ -190,7 +196,7 @@ constexpr bool EpArgsOffsetsAscend() {
 
 }  // namespace detail
 
-static_assert(detail::kEpArgsFieldCount == 22,
+static_assert(detail::kEpArgsFieldCount == 24,
               "added an EpArgs field -- add it to MORI_EP_ARGS_FIELDS in the same position "
               "and bump this count");
 static_assert(detail::EpArgsOffsetsAscend(),
@@ -218,6 +224,10 @@ struct EpCfg {
 
   // ---- algorithm ----
   bool useWeights = true;
+  // Per-token scale row carried alongside the payload, in bytes. 0 = off, and off
+  // is free: Render omits default-valued fields, so the Cfg text -- which IS the
+  // JIT cache key -- is byte-identical to a build without this feature.
+  int scaleBytes = 0;
 };
 
 template <typename Self, typename Visit>
@@ -234,10 +244,11 @@ inline void VisitFields(Self& c, const EpCfg& d, Visit&& v) {
   MORI_FIELD(warpPerBlock);
   MORI_FIELD(waveSize);
   MORI_FIELD(useWeights);
+  MORI_FIELD(scaleBytes);
 #undef MORI_FIELD
 }
 
-MORI_JIT_ASSERT_FIELD_COUNT(EpCfg, 11, "added an EpCfg field -- update VisitFields(EpCfg) too");
+MORI_JIT_ASSERT_FIELD_COUNT(EpCfg, 12, "added an EpCfg field -- update VisitFields(EpCfg) too");
 
 inline std::string Render(const EpCfg& c) {
   const EpCfg d{};
@@ -305,16 +316,64 @@ constexpr int EpCombineSharedBytes(const EpCfg& c) {
 
 constexpr int EpTokenBytes(const EpCfg& c) { return c.hiddenDim * EpElemSize(c.dtype); }
 
+// The scale row's SLOT stride: the caller's row padded to 128 B. A transfer is a
+// run of consecutive slots, and TdmWholeOrSplit128 only gives a body to the part
+// of a run that starts aligned -- at the natural 224 B only every 4th one does.
+// gfx1250, 512 tok/rank, hidden 7168, EP4: dispatch 157.7us at 224 B, 94.8 at 256.
+//
+// Padded on every arch so the destination layout does not fork per arch. The pad
+// is unconditional, so small rows amplify: 224->256 is 1.14x, 32->128 is 4x,
+// 4->128 is 32x, and the gfx1250 staging pool is worldSize^2 * maxTok * THIS per
+// compiled variant. Bounded by the static_assert there, not by this being cheap.
+//
+// Everything reading the destination uses this, not Cfg.scaleBytes; it assumes
+// the arena aligns the region too (SymmArena._ALIGN, checked in hip_backend).
+constexpr int EpScaleAlign = 128;
+constexpr int EpScaleStride(const EpCfg& c) {
+  return c.scaleBytes <= 0 ? 0 : (c.scaleBytes + EpScaleAlign - 1) / EpScaleAlign * EpScaleAlign;
+}
+
 // gfx1250 launch LDS. Dispatch stages one token tile per warp through the TDM
 // engine; combine reserves the whole budget and sizes its tiles at runtime.
 // EpCombine1250xLdsBudget must match MORI_COMB_LDS_BUDGET in ep_intranode_1250x.hpp.
-constexpr int EpCombine1250xLdsBudget = 327680;
+//
+// Ep1250xLdsBytes is the physical ceiling; combine's budget happens to be all of
+// it. Anything else that needs the ceiling should say so directly rather than
+// reach for combine's number -- otherwise retuning combine silently resizes it.
+constexpr int Ep1250xLdsBytes = 327680;
+constexpr int EpCombine1250xLdsBudget = Ep1250xLdsBytes;
 // xdbFlag slots for the per-block combine entry barrier: one uint64 per block, so a
 // block is the only writer of its own epoch. 256 == the CU count, which caps the
 // combine block_num; the host allocates this many and every call keeps them in step.
 constexpr int EpXdbFlagSlots = 256;
+
+// Per-warp LDS slab. The metadata tile and the payload tile share it (same
+// address, different phases), so its size bounds BOTH -- and the metadata batch
+// size is (slab - headroom) / bytes-per-token.
+//
+// Sizing it off the payload dtype alone is fine until a scale row joins the
+// metadata: an fp8 payload halves the slab exactly when per-token metadata grows
+// from 52 to 276 bytes, and the batch collapses ~11x.
+//
+// The floor is EMPIRICAL, and from one shape: hidden 7168, topk 6, EP4, 512
+// tokens/rank on gfx1250, where it moved dispatch from 93.1us back to 36.7us.
+// The mechanism (a narrower slab shrinks the metadata batch) is general; the
+// conclusion that the bf16 width is the right target is not necessarily so for a
+// very different hidden size, and this is the first thing to re-measure if one
+// shows up.
+//
+// Capped by the physical LDS rather than by combine's budget: the two are the
+// same number today, but for a shared physical reason, not because dispatch
+// follows combine.
+constexpr int EpDispatch1250xSlabBytes(const EpCfg& c) {
+  const int payload = c.hiddenDim * EpElemSize(c.dtype);
+  if (c.scaleBytes <= 0) return payload;
+  const int wide = c.hiddenDim * EpElemSize(EpDType::Bf16);
+  const long long total = (long long)wide * c.warpPerBlock;
+  return (wide > payload && total <= Ep1250xLdsBytes) ? wide : payload;
+}
 constexpr int EpDispatch1250xLdsBytes(const EpCfg& c) {
-  return c.warpPerBlock * c.hiddenDim * EpElemSize(c.dtype);
+  return c.warpPerBlock * EpDispatch1250xSlabBytes(c);
 }
 
 // A Cfg that cannot launch is a host-side error, not a kernel that misbehaves.
@@ -332,7 +391,9 @@ constexpr bool EpCfgIsValid(const EpCfg& c) {
          // per peer / per top-k slot within a single wavefront.
          c.numExpertPerToken < c.waveSize && c.worldSize <= c.waveSize &&
          // WarpCopy moves whole 16 B chunks.
-         (EpTokenBytes(c) % 16) == 0;
+         (EpTokenBytes(c) % 16) == 0 &&
+         // Scale rows are copied as dwords, so the row must be dword-sized.
+         c.scaleBytes >= 0 && (c.scaleBytes % 4) == 0;
 }
 
 }  // namespace v2

include/mori/ops/dispatch_combine_v2/ep_spec.hpp

@@ -50,6 +50,7 @@ struct EpRequest {
   bool useWeights = true;
   int blockNum = 0;      // 0 = arch default
   int warpPerBlock = 0;  // 0 = arch default
+  int scaleBytes = 0;    // per-token scale row carried with the payload; 0 = off
 };
 
 template <typename Self, typename Visit>
@@ -65,10 +66,11 @@ inline void VisitFields(Self& r, const EpRequest& d, Visit&& v) {
   MORI_FIELD(useWeights);
   MORI_FIELD(blockNum);
   MORI_FIELD(warpPerBlock);
+  MORI_FIELD(scaleBytes);
 #undef MORI_FIELD
 }
 
-MORI_JIT_ASSERT_FIELD_COUNT(EpRequest, 10,
+MORI_JIT_ASSERT_FIELD_COUNT(EpRequest, 11,
                             "added an EpRequest field -- update VisitFields(EpRequest) too");
 
 std::string EpRequestSchema();

python/mori/jit/v2/plan_api.py

@@ -64,6 +64,16 @@
 # (the caller halves hiddenDim). Combine cannot use it; C++ rejects that.
 DTYPES = {"bf16": 0, "fp32": 1, "byte8": 2}
 
+# Arena regions a caller is allowed not to carry. Everything else missing is a
+# bug in the caller, not a configuration: see the bind loop in make_plan.
+# Keyed by the snake_case region name: the region key comes from the C++ launch
+# argument (offOutScales -> "outScales") while an arena names it "out_scales",
+# so both are folded through _camel_to_snake before being looked up here.
+_OPTIONAL_REGIONS = frozenset({"out_scales"})
+
+# ... and the Request field that makes each of them mandatory again.
+_REGION_REQUIRED_WHEN = {"out_scales": "scaleBytes"}
+
 # The C ABI + the plan registry both live here; op-libraries register INTO it.
 _ABI_NAME = "libmori_jit.so"
 
@@ -418,7 +428,34 @@ def __init__(self, *, arch=None, arena=None, region_names=None, **kwargs):
             if arena is not None:
                 names = region_names or {}
                 for wire, region in arg_regions.items():
-                    self._defaults[wire] = int(arena.offset(names.get(region, region)))
+                    name = names.get(region, region)
+                    try:
+                        self._defaults[wire] = int(arena.offset(name))
+                    except KeyError:
+                        # ONLY a region on the optional list may be missing. A
+                        # blanket catch here would turn a typo in region_names, or
+                        # a backend that forgot a required region, from an
+                        # immediate KeyError into a silent bind to offset 0 -- and
+                        # 0 aliases the first region, so the kernel would scribble
+                        # over it. Callers that build their own arena and pass a
+                        # region_names mapping (aiter's MegaMoE does) are exactly
+                        # the ones most able to get this wrong.
+                        canon = _camel_to_snake(region)
+                        if canon not in _OPTIONAL_REGIONS:
+                            raise
+                        # Optional, but not optional for THIS plan: if the Request
+                        # field that turns the feature on is set, the kernel will
+                        # dereference the offset and a 0 would alias region 0.
+                        # Mechanism, not a comment telling callers to be careful.
+                        enabler = _REGION_REQUIRED_WHEN.get(canon)
+                        if enabler and int(req.get(enabler) or 0):
+                            raise KeyError(
+                                f"{kernel}: arena has no region {name!r}, but "
+                                f"{enabler}={req[enabler]} turns it on"
+                            ) from None
+                        # Optional and genuinely off: the kernel's `if constexpr`
+                        # is what keeps the 0 from being dereferenced.
+                        self._defaults[wire] = 0
 
         def bind(self, **args) -> None:
             """Pin launch arguments that never change between calls.
@@ -472,9 +509,7 @@ def launch(self, stream=0, **args) -> None:
                         setattr(buf, name, _as_ptr(merged[name]))
                 self._buf = buf
                 self._buf_shape = shape
-                self._dyn_args = tuple(
-                    (k, arg_snake_to_wire.get(k, k)) for k in args
-                )
+                self._dyn_args = tuple((k, arg_snake_to_wire.get(k, k)) for k in args)
                 self._dyn_defs = tuple(
                     w
                     for w, v in self._defaults.items()

python/mori/ops/dispatch_combine_v2/dispatch_combine_op.py

@@ -575,6 +575,20 @@ def _gate(self, kernels: KernelSet) -> None:
     def capabilities(self) -> frozenset[str]:
         return self._kernels.capabilities
 
+    def scale_stride_bytes(self) -> int:
+        """Bytes between consecutive out_scales rows, 0 when off.
+
+        A backend may lay them down wider than the row it was handed (the HIP one
+        pads to 128 B). On the base so a consumer never branches on the backend;
+        the default is the row unchanged, which is right for one that does not pad.
+        """
+        return self._scale_row_bytes()
+
+    def _scale_row_bytes(self) -> int:
+        """The caller's row in bytes, dword-rounded. 0 when the transport is off."""
+        n = self.cfg.scale_dim * self.cfg.scale_type_size
+        return ((n + 3) // 4) * 4 if n else 0
+
     # -- shared: variant selection -----------------------------------------
 
     def _pick(self, num_tokens):
@@ -667,6 +681,10 @@ def dispatch(
         total_recv[, routing]); out == arena disp_out, safe to read without
         .clone() because combine stages into a separate out_tok buffer.
 
+        out_scales is [max_recv, scale_dim_i32] on both backends but is not always
+        CONTIGUOUS -- a backend may lay the rows down wider and return a strided
+        view. Read by pointer, stride by scale_stride_bytes(), not by this shape.
+
         total_recv is a DEVICE tensor. Reading it on the host is a full sync that
         costs more than the kernel and makes the op uncapturable, so neither
         backend does it.