track1(batch4): thread micro-batch axis end-to-end through path_c gridded step

Thread batch=N (bs1 default, bs4 selectable) through the path_c shape-config ->
delegation prim -> physical ABI bank sizing -> probe, keeping bs1 byte-identical.

- PathCModelShapeEnv: add batch:int=1 field + __post_init__ RAISE on batch<1
  (no silent clamp); _path_c_model_shape_env_from_config reads config.micro_batch_size
  (RAISE if present-but-non-positive); replace() preserves batch.
- ModelFactoryProfile: add micro_batch_size:int=1 (+_require_positive); hybrid_config()
  carries it onto the frozen HybridTinyConfig via object.__setattr__ as the single
  source of truth, reads it back + RAISES if it failed to stick (no bsN->bs1 revert).
- delegation prim call site: pass batch=int(resolved_shape_env.batch) (kills the
  silent batch=1 default trap) + add a RULE #1 batch-consistency guard that RAISEs
  with buffer-name + expected-vs-got batch if a compiled grid kernel's leading dim
  disagrees with the region-declared batch.
- _mamba3_chunked_fwd_region_buffer_shape: batch=int(shape_env.batch) (resizes ALL
  chunked region/handoff banks); serial-emit feasibility threads shape_env.batch.
- _buffer_shape: apply batch multiplier EXPLICITLY, per-buffer, ONLY to per-token
  activation/state/checkpoint banks (hidden, mamba/m2rnn state+checkpoint+conv_state,
  target_ids/mask, q_fp8/kv_fp8/scales/indices/lse, *_hidden/_delta/_out); weights
  stay 1x (verified: act/state scale 4x, weight banks unchanged).
- path_c_relax_step_banks: add real_bank_numels_for_batch() + assert_batch_scaling()
  bs4 self-check (param/paramg banks must stay 1x; act/actg/state must grow -> RAISE).
- probes: add --bs4 (requires --prod; RAISE otherwise) flipping the prod cfg batch
  1->4; fwd probe adds an F0/F2 grid-total == 4x bs1 sanity assert. Builders/grids
  already consume batch; no kernel edit.

RULE #1: every batch path RAISES on mismatch; no bs4->bs1 fallback anywhere.
Local check: ast.parse + py_compile (py3.13) PASS on all 6; isolated batch-logic +
per-buffer scaling discipline validated (gb10 compile/run is the Profile agent).

Author: apstenku123

cppmega_mlx/recipes/model_factory.py

@@ -132,6 +132,13 @@ class ModelFactoryProfile:
     num_attention_heads: int
     head_dim: int
     vocab_size: int
+    # Outer micro-batch axis for the path_c gridded step. Default 1 keeps every
+    # existing profile/consumer at the historical bs1 behavior byte-for-byte; set
+    # it to 4 (via ``local_gb10_quarter_profile(micro_batch_size=4)`` or a
+    # ``replace``) to select the bs4 gridded path. It is threaded onto the
+    # HybridTinyConfig the profile yields so the path_c shape-env reads it as the
+    # single source of truth. RULE #1: a non-positive value RAISES in __post_init__.
+    micro_batch_size: int = 1
     max_seq_length: int = LOCAL_GB10_QUARTER_MAX_SEQ_LENGTH
     dsa_a_layer_ranks: tuple[int, ...] = LOCAL_GB10_QUARTER_DSA_A_LAYER_RANKS
     dsa_indexer_n_heads: int | None = 32
@@ -161,6 +168,7 @@ def __post_init__(self) -> None:
         _require_positive("num_attention_heads", self.num_attention_heads)
         _require_positive("head_dim", self.head_dim)
         _require_positive("vocab_size", self.vocab_size)
+        _require_positive("micro_batch_size", self.micro_batch_size)
         _require_positive("max_seq_length", self.max_seq_length)
         _require_positive("moe_num_experts", self.moe_num_experts)
         _require_positive("moe_top_k", self.moe_top_k)
@@ -217,9 +225,43 @@ def nam56r_config(self) -> Nam56RModelConfig:
         )
 
     def hybrid_config(self, **overrides) -> HybridTinyConfig:
-        """Return a HybridTinyConfig using existing NAM56R-to-MLX mapping."""
-
-        return build_hybrid_tiny_config_from_nam56r(self.nam56r_config(), **overrides)
+        """Return a HybridTinyConfig using existing NAM56R-to-MLX mapping.
+
+        The profile's ``micro_batch_size`` rides onto the returned config as the
+        single source of truth for the path_c gridded step's outer batch axis
+        (``_path_c_model_shape_env_from_config`` reads ``config.micro_batch_size``).
+        ``HybridTinyConfig`` is a frozen dataclass with no such field, so the value
+        is carried directly on the instance via ``object.__setattr__`` (the
+        documented "carry it on the config/region directly" path). RULE #1: an
+        explicit ``micro_batch_size`` override that conflicts with the profile's is
+        rejected, and the carried value is read back and asserted — a value that
+        failed to stick RAISES rather than silently reverting to bs1.
+        """
+
+        requested_batch = int(self.micro_batch_size)
+        if "micro_batch_size" in overrides:
+            override_batch = int(overrides.pop("micro_batch_size"))
+            if override_batch < 1:
+                raise ValueError(
+                    "hybrid_config(micro_batch_size=...) must be a positive int "
+                    f"(got {override_batch!r}); RULE #1 forbids a silent clamp to 1"
+                )
+            requested_batch = override_batch
+        config = build_hybrid_tiny_config_from_nam56r(
+            self.nam56r_config(), **overrides
+        )
+        # Carry the micro-batch axis onto the frozen config instance (no field in
+        # HybridTinyConfig) so the path_c shape-env reads ONE source of truth.
+        object.__setattr__(config, "micro_batch_size", requested_batch)
+        carried = int(getattr(config, "micro_batch_size", 0))
+        if carried != requested_batch:
+            raise RuntimeError(
+                "model_factory.hybrid_config: failed to carry micro_batch_size "
+                f"onto the HybridTinyConfig (requested {requested_batch}, read back "
+                f"{carried!r}); RULE #1: refusing a silent bs{requested_batch}->bs1 "
+                "revert"
+            )
+        return config
 
     def tiny_smoke_config(self, **overrides) -> HybridTinyConfig:
         """Return a small T=512-capable config preserving this route profile."""

cppmega_mlx/runtime/path_c_fusion.py

@@ -433,6 +433,19 @@ class PathCModelShapeEnv:
     m2rnn_conv_kernel: int
     model_value_dtype: str = "float32"
     vocab_size: int = 0
+    # Outer micro-batch axis threaded into every gridded F0/F1/F2/B0/B1/B2 builder
+    # and into the per-token activation/state bank shapes. Default 1 keeps the
+    # historical bs1 behavior byte-identical for every existing caller; a >1 value
+    # selects the bs4 (or any bsN) gridded path. RULE #1: an invalid batch RAISES
+    # in __post_init__ (no silent clamp to 1).
+    batch: int = 1
+
+    def __post_init__(self) -> None:
+        if int(self.batch) < 1:
+            raise ValueError(
+                "PathCModelShapeEnv.batch must be a positive micro-batch size "
+                f"(got {self.batch!r}); RULE #1 forbids a silent clamp to 1"
+            )
 
     @property
     def mamba_inner_dim(self) -> int:
@@ -1438,9 +1451,26 @@ def _path_c_model_shape_env_from_config(config: Any | None) -> PathCModelShapeEn
         hidden_size = int(getattr(config, "hidden_size"))
     except AttributeError:
         return None
+    # Micro-batch axis (bs1 default for every config that does not carry it; bs4
+    # and friends ride on ``micro_batch_size``). RULE #1: a config that DOES carry
+    # a ``micro_batch_size`` must yield a usable positive int — if it is present
+    # but non-coercible/non-positive we RAISE rather than silently defaulting to 1
+    # (a silent bs4->bs1 revert is exactly the forbidden fallback).
+    raw_batch = getattr(config, "micro_batch_size", None)
+    if raw_batch is None:
+        batch = 1
+    else:
+        batch = int(raw_batch)
+        if batch < 1:
+            raise ValueError(
+                "Path C model shape-env: config.micro_batch_size must be a "
+                f"positive int (got {raw_batch!r}); RULE #1 forbids a silent "
+                "clamp/revert to bs1"
+            )
     return PathCModelShapeEnv(
         sequence_length=sequence_length,
         hidden_size=hidden_size,
+        batch=batch,
         attention_num_q_heads=int(attention_config.num_q_heads),
         attention_num_kv_heads=int(attention_config.kv_heads),
         attention_head_dim=int(attention_config.q_head_dim),

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -2377,6 +2377,13 @@ def build_path_c_descriptor_prim_func(
             op_name=nodes[0].op_name,
             production_source=descriptors[0].production_source,
             shape_env=resolved_shape_env,
+            # Thread the micro-batch axis EXPLICITLY: the delegation prim's
+            # ``batch: int = 1`` default is the canonical silent-fallback trap —
+            # if we passed no batch= the F0..B2 grid kernels would compile bs1
+            # even when the surrounding step is bs4 (degraded/wrong, not a crash).
+            # Sourcing it from ``resolved_shape_env.batch`` keeps banks and kernels
+            # on the SAME single source of truth (RULE #1).
+            batch=int(resolved_shape_env.batch),
         )
         delegated_prim._cppmega_path_c_schedule_generator = (
             PATH_C_DESCRIPTOR_SCHEDULE_GENERATOR
@@ -2430,6 +2437,49 @@ def build_path_c_descriptor_prim_func(
             for param in tuple(getattr(delegated_prim, "params", ()))
             if not (hasattr(param, "is_scalar") and bool(param.is_scalar()))
         )
+        # RULE #1 BATCH-CONSISTENCY GUARD (no silent bs4->bs1): the region surface
+        # sized its handoff/region banks for ``resolved_shape_env.batch`` (via
+        # ``_mamba3_chunked_fwd_region_buffer_shape``); the compiled grid kernel was
+        # compiled for the SAME batch (passed above). If a half-threaded path let
+        # them disagree, a bs4 kernel writing into a bs1-sized bank would be a
+        # silent out-of-bounds / truncation. Cross-check the compiled kernel's
+        # device-buffer param leading dim against the region-declared batch for
+        # every buffer whose region shape carries an explicit leading ``batch``
+        # axis (the per-token/per-state chunked banks; the bare ``A``/``D`` vectors
+        # and any scalar slots carry no batch axis and are skipped). RAISE with
+        # WHERE (buffer name) + WHAT (expected batch vs prim leading dim).
+        _expected_batch = int(resolved_shape_env.batch)
+        if len(delegated_buffer_order) == len(delegated_buffer_params):
+            for _bufname, _param in zip(
+                delegated_buffer_order, delegated_buffer_params
+            ):
+                _region_shape = _mamba3_chunked_fwd_region_buffer_shape(
+                    _bufname, resolved_shape_env
+                )
+                if _region_shape is None or len(_region_shape) == 0:
+                    continue
+                # The chunked region buffers that carry an outer batch axis declare
+                # it as dim 0 == batch (x/B/C/dt/h0/cb/.../final_state/d*). Vectors
+                # like A=(nheads,)/D=(nheads,) have a leading dim != batch by design
+                # — skip those (their dim 0 is a head/feature count, not batch).
+                if int(_region_shape[0]) != _expected_batch:
+                    continue
+                _prim_shape = tuple(
+                    int(dim) for dim in tuple(getattr(_param, "shape", ()))
+                )
+                if len(_prim_shape) == 0:
+                    continue
+                if int(_prim_shape[0]) != _expected_batch:
+                    raise ValueError(
+                        "mamba3 chunked-scan delegation batch mismatch for buffer "
+                        f"{_bufname!r}: region surface sized leading dim "
+                        f"{_expected_batch} (shape_env.batch) but the compiled grid "
+                        f"kernel param leading dim is {int(_prim_shape[0])} "
+                        f"(prim shape {_prim_shape}). A half-threaded bs"
+                        f"{_expected_batch} path would write OOB into a bs"
+                        f"{int(_prim_shape[0])}-sized bank (RULE #1: no silent "
+                        "bs4->bs1 fallback)."
+                    )
         # Attach the named-buffer ABI only when the region surface's ordered
         # buffers RECONCILE 1:1 with the kernel's device-buffer param slots. The
         # FULL model surface declares exactly the kernel's inputs+outputs (verified
@@ -7499,7 +7549,9 @@ def _append_row_phased_mamba3_body(
             _chunked_limit,
         ) = _mamba3_chunked_forward_scan_feasibility(
             sequence_length=int(shape_env.sequence_length),
-            batch=1,
+            # Thread the micro-batch axis for consistency (flag-OFF serial-emit
+            # classification only; does not change grid feasibility True/False).
+            batch=int(shape_env.batch),
             heads=heads,
             head_dim=head_dim,
             state_dim=state_dim,
@@ -14691,7 +14743,11 @@ def _mamba3_chunked_fwd_region_buffer_shape(
             break
     if suffix is None:
         return None
-    batch = 1
+    # Outer micro-batch axis: every chunked region/handoff bank below carries
+    # ``batch`` as leading dim, so this single source resizes ALL of them to bs4
+    # (or any bsN). RULE #1: sourced from the SAME shape_env the delegation prim
+    # compiles against, so banks and kernels cannot disagree.
+    batch = int(shape_env.batch)
     seqlen = int(shape_env.sequence_length)
     nheads = int(shape_env.mamba_num_heads)
     headdim = int(shape_env.mamba_head_dim)
@@ -14755,6 +14811,15 @@ def _buffer_shape(
     q_dim = q_heads * head_dim
     kv_dim = kv_heads * head_dim
     topk = shape_env.attention_sparse_topk
+    # Outer micro-batch axis. RULE #1: applied EXPLICITLY, per-buffer, ONLY to the
+    # per-token activation/state/checkpoint banks below — NEVER to weight banks
+    # (mamba3_in_proj_weight, *_norm_weight, lm_head_weight, dt_bias, D,
+    # conv_weight, ...), which are batch-INVARIANT. A blanket multiply would 4x the
+    # parameter banks (wrong) and break the param-1x liveness model + the bs1
+    # self-check. The chunked region buffers (x/B/C/cb/prev_states/...) already
+    # carried batch via _mamba3_chunked_fwd_region_buffer_shape above and never
+    # reach here.
+    batch = int(shape_env.batch)
     if name in {
         "hidden",
         "mamba3_delta",
@@ -14765,28 +14830,36 @@ def _buffer_shape(
         "hidden_after_m2rnn",
         "attention_out",
     }:
-        return (seq * hidden,)
+        return (batch * seq * hidden,)
     if _is_mamba3_state_like_buffer(buffer_name, name):
         return (
-            shape_env.mamba_num_heads
+            batch
+            * shape_env.mamba_num_heads
             * shape_env.mamba_head_dim
             * shape_env.mamba_state_dim,
         )
     if name == "mamba3_angle_state":
         return (
-            max(1, shape_env.mamba_num_heads * shape_env.mamba_num_rope_angles),
+            max(
+                1,
+                batch * shape_env.mamba_num_heads * shape_env.mamba_num_rope_angles,
+            ),
         )
     if name == "mamba3_angle_grad_state":
         return (
-            max(1, shape_env.mamba_num_heads * shape_env.mamba_num_rope_angles),
+            max(
+                1,
+                batch * shape_env.mamba_num_heads * shape_env.mamba_num_rope_angles,
+            ),
         )
     if name == "mamba3_angle_checkpoint":
         checkpoint_interval = MAMBA3_BWD_REPLAY_CHECKPOINT_INTERVAL
         checkpoint_count = (
             shape_env.sequence_length + checkpoint_interval - 1
         ) // checkpoint_interval
         return (
-            (checkpoint_count + 1)
+            batch
+            * (checkpoint_count + 1)
             * max(1, shape_env.mamba_num_heads * shape_env.mamba_num_rope_angles),
         )
     if name == "mamba3_h_checkpoint":
@@ -14795,7 +14868,8 @@ def _buffer_shape(
             shape_env.sequence_length + checkpoint_interval - 1
         ) // checkpoint_interval
         return (
-            (checkpoint_count + 1)
+            batch
+            * (checkpoint_count + 1)
             * shape_env.mamba_num_heads
             * shape_env.mamba_head_dim
             * shape_env.mamba_state_dim,
@@ -14804,7 +14878,9 @@ def _buffer_shape(
         return (
             max(
                 1,
-                (shape_env.mamba_conv_kernel - 1) * shape_env.mamba_conv_channels,
+                batch
+                * (shape_env.mamba_conv_kernel - 1)
+                * shape_env.mamba_conv_channels,
             ),
         )
     if name == "mamba3_in_proj_weight":
@@ -14841,7 +14917,7 @@ def _buffer_shape(
         # the eager ``layers.{first_in_region}.norm.weight`` parameter.
         return (hidden,)
     if name in {"target_ids", "target_mask"}:
-        return (seq,)
+        return (batch * seq,)
     if name == "lm_head_weight":
         vocab = max(1, int(getattr(shape_env, "vocab_size", 0) or 0))
         return (vocab * hidden,)
@@ -14865,7 +14941,8 @@ def _buffer_shape(
         return (hidden * shape_env.m2rnn_num_heads * shape_env.m2rnn_v_head_dim,)
     if name in {"m2rnn_h0", "m2rnn_h_state"}:
         return (
-            shape_env.m2rnn_num_heads
+            batch
+            * shape_env.m2rnn_num_heads
             * shape_env.m2rnn_k_head_dim
             * shape_env.m2rnn_v_head_dim,
         )
@@ -14875,13 +14952,16 @@ def _buffer_shape(
             shape_env.sequence_length + checkpoint_interval - 1
         ) // checkpoint_interval
         return (
-            (checkpoint_count + 1)
+            batch
+            * (checkpoint_count + 1)
             * shape_env.m2rnn_num_heads
             * shape_env.m2rnn_k_head_dim
             * shape_env.m2rnn_v_head_dim,
         )
     if name == "m2rnn_conv_state":
-        return ((shape_env.m2rnn_conv_kernel - 1) * shape_env.m2rnn_conv_dim,)
+        return (
+            batch * (shape_env.m2rnn_conv_kernel - 1) * shape_env.m2rnn_conv_dim,
+        )
     if name == "attention_q_proj_weight":
         return (q_dim * hidden,)
     if name == "attention_q_proj_bias":
@@ -14897,17 +14977,17 @@ def _buffer_shape(
     if name == "attention_out_proj_bias":
         return (hidden,)
     if name == "q_fp8":
-        return (seq * q_heads * head_dim,)
+        return (batch * seq * q_heads * head_dim,)
     if name == "q_scale":
-        return (seq * q_heads,)
+        return (batch * seq * q_heads,)
     if name == "kv_fp8":
-        return (seq * kv_heads * head_dim,)
+        return (batch * seq * kv_heads * head_dim,)
     if name == "kv_scale":
-        return (seq * kv_heads,)
+        return (batch * seq * kv_heads,)
     if name == "indices":
-        return (seq * kv_heads * topk,)
+        return (batch * seq * kv_heads * topk,)
     if name == "lse":
-        return (seq * q_heads,)
+        return (batch * seq * q_heads,)
     if name == "sparse_mla_sm_scale":
         return (1,)
     if name == "sparse_mla_sinks":
@@ -14920,7 +15000,7 @@ def _buffer_shape(
         or name.endswith("_delta")
         or name.endswith("_out")
     ):
-        return (seq * hidden,)
+        return (batch * seq * hidden,)
     return (buffer_extent,)