feat(v4/fusion): Stage B — descriptor synthesizer for V4 bricks

Adds cppmega_v4.fusion.descriptor_synthesizer covering every kind in
BLOCK_BUILDERS with a structurally valid PathCBrickScheduleDescriptor:

  - synthesize_descriptor_for_brick(kind) builds a fallback descriptor
    whose schedule_family / required_codegen_steps / supports_backward
    are derived from the brick's FusionEligibility category. Opaque
    bricks (SDPA, cross-attn, sparse, recurrence) are tagged
    implementation_status="opaque_brick_passthrough" so downstream
    consumers group them into regions via DLPack handoff instead of
    inlining them into a fused PrimFunc.

  - build_v4_extended_registry() returns a registry that combines the
    legacy v3 descriptors with v4 fallbacks; legacy entries always
    win on collision (descriptor_for never returns the synth when a
    hand-written one exists).

This unblocks descriptors_for_signature() for any chain assembled from
BLOCK_BUILDERS so Stage C's auto-planner can call Path C codegen
without None lookups.

Tests (tests/v4/test_fusion_stage_b.py, 28 cases):
  - per-brick synth: known kinds (gdn linear_attn, gated_attention
    opaque sdpa), unknown kind safe passthrough
  - parametrized: every BLOCK_BUILDERS kind produces a
    registry-acceptable descriptor
  - registry: legacy preserved, no clobber, full kind coverage,
    AOT <op>_bwd synth fires only for supports_backward=True
  - system: signature lookup for Qwen3-Next (3xGDN + gated + moe) and
    Ling 2.6 (7x bailing_linear + mla + moe) chains

Full v4 regression: 450 passed / 5 skipped / 0 failed.

Author: apstenku123

cppmega_v4/fusion/__init__.py

@@ -27,6 +27,10 @@
     FusionEligibility,
     can_fuse_pair,
 )
+from cppmega_v4.fusion.descriptor_synthesizer import (
+    build_v4_extended_registry,
+    synthesize_descriptor_for_brick,
+)
 from cppmega_v4.fusion.dlpack_bridge import (
     dlpack_available,
     host_copy_fallback,
@@ -38,11 +42,13 @@
     "BrickGraph",
     "BrickNode",
     "FusionEligibility",
+    "build_v4_extended_registry",
     "can_fuse_pair",
     "dlpack_available",
     "from_block_specs",
     "from_mlx_model",
     "host_copy_fallback",
     "mlx_to_tilelang",
+    "synthesize_descriptor_for_brick",
     "tilelang_to_mlx",
 ]

cppmega_v4/fusion/descriptor_synthesizer.py

@@ -0,0 +1,170 @@
+"""Stage B — auto-synthesize Path C brick descriptors for V4 bricks.
+
+Background:
+  ``cppmega_mlx.runtime.path_c_fusion_schedules.default_path_c_brick_schedule_descriptor_registry``
+  ships 5 hand-written descriptors (legacy v3 train block). The 15+ V4
+  bricks in ``cppmega_v4.models.unified_superblock_v4.BLOCK_BUILDERS``
+  have no descriptors, so any region built from them fails the
+  ``descriptors_for_signature`` lookup and Path C codegen bails.
+
+Strategy:
+  - For each V4 brick, synthesise a "fallback descriptor" that marks the
+    op as an opaque kernel (already implemented in Path B / mlx-lm SDPA /
+    upstream) but is structurally valid for the registry: non-empty
+    op_name, well-formed required_codegen_steps, schedule_family chosen
+    by the FusionEligibility category (linear_attn / sdpa_attention /
+    moe / ssm / nonlinear_rnn / cross_attn / sparse_attn / norm_or_proj).
+  - The fallback descriptor's ``implementation_status`` is
+    ``"opaque_brick_passthrough"`` — that signals to downstream pattern
+    matchers that the brick must NOT be inlined into a fused PrimFunc,
+    but it CAN be grouped into a region (and the planner can lower
+    surrounding bricks fused with DLPack hand-off on the brick's edges).
+
+Public surface:
+  - ``synthesize_descriptor_for_brick(kind) -> PathCBrickScheduleDescriptor``
+  - ``build_v4_extended_registry() -> PathCBrickScheduleDescriptorRegistry``
+    Combines the default v3 registry with v4 fallback descriptors. All
+    BLOCK_BUILDERS kinds get a descriptor; existing v3 entries are kept
+    as-is (no clobber).
+"""
+
+from __future__ import annotations
+
+from cppmega_mlx.runtime.path_c_fusion_schedules import (
+    PathCBrickScheduleDescriptor,
+    PathCBrickScheduleDescriptorRegistry,
+    default_path_c_brick_schedule_descriptor_registry,
+)
+from cppmega_v4.fusion.compatibility import _CATEGORY_BY_KIND
+from cppmega_v4.models.unified_superblock_v4 import BLOCK_BUILDERS
+
+
+# Map fusion category -> schedule_family token (descriptor metadata).
+# Categories that fuse via path_c get a family hint matching the kind of
+# loop they're built around; opaque ones get a sentinel that downstream
+# pattern matchers treat as "do not inline".
+_FAMILY_BY_CATEGORY: dict[str, str] = {
+    "linear_attn": "linear_attn_scan",
+    "ssm": "ssm_chunkwise_scan",
+    "sdpa_attention": "opaque_sdpa_with_outputs",
+    "cross_attn": "opaque_cross_attn",
+    "nonlinear_rnn": "opaque_recurrence",
+    "moe": "moe_route_combine",
+    "sparse_attn": "opaque_sparse",
+    "norm_or_proj": "pointwise_or_linear",
+    "unknown": "opaque_passthrough",
+}
+
+
+def _required_codegen_steps_for(kind: str, category: str) -> tuple[str, ...]:
+    """Per-brick codegen step sequence. Used by descriptor consumers to
+    decide which sub-templates to thread together. Conservative defaults
+    so an unknown brick still passes the registry's structural check."""
+    if category == "linear_attn":
+        return (
+            f"{kind}_scan_descriptor",
+            f"{kind}_state_init",
+        )
+    if category == "ssm":
+        return (
+            f"{kind}_ssm_descriptor",
+            f"{kind}_chunkwise_internal_buffers",
+        )
+    if category == "sdpa_attention":
+        return (
+            f"{kind}_opaque_sdpa_descriptor",
+            f"{kind}_output_gate_o_proj_fragment",
+        )
+    if category == "cross_attn":
+        return (
+            f"{kind}_opaque_cross_attn_descriptor",
+            f"{kind}_residual_output_fragment",
+        )
+    if category == "nonlinear_rnn":
+        return (
+            f"{kind}_opaque_recurrence_descriptor",
+            f"{kind}_pre_norm_fragment",
+        )
+    if category == "moe":
+        return (
+            f"{kind}_route_descriptor",
+            f"{kind}_expert_ffn_fragment",
+            f"{kind}_combine_descriptor",
+        )
+    if category == "sparse_attn":
+        return (
+            f"{kind}_opaque_sparse_descriptor",
+        )
+    if category == "norm_or_proj":
+        return (
+            f"{kind}_pointwise_descriptor",
+        )
+    return (f"{kind}_opaque_passthrough",)
+
+
+def _supports_backward(category: str) -> bool:
+    # Conservative: only the categories with established gradient paths
+    # in our codebase claim backward support. SDPA-backed and sparse
+    # bricks have backwards via mlx-lm/MLX autodiff but not via Path C
+    # AOT — the descriptor itself doesn't know how to emit a backward
+    # PrimFunc fragment, so it sets supports_backward=False.
+    return category in {"linear_attn", "ssm", "norm_or_proj", "nonlinear_rnn"}
+
+
+def synthesize_descriptor_for_brick(kind: str) -> PathCBrickScheduleDescriptor:
+    """Build a fallback descriptor for a single brick kind.
+
+    Returns a structurally valid PathCBrickScheduleDescriptor that the
+    registry accepts and that downstream consumers can treat as an
+    opaque-brick handoff. Never raises for unknown kinds — falls back to
+    ``opaque_passthrough`` family.
+    """
+    category = _CATEGORY_BY_KIND.get(kind, "unknown")
+    family = _FAMILY_BY_CATEGORY.get(category, "opaque_passthrough")
+    is_opaque = family.startswith("opaque_")
+    return PathCBrickScheduleDescriptor(
+        op_name=kind,
+        implementation_status=(
+            "opaque_brick_passthrough" if is_opaque else "descriptor_synth_v4"
+        ),
+        required_codegen_steps=_required_codegen_steps_for(kind, category),
+        schedule_family=family,
+        supports_backward=_supports_backward(category),
+        description=(
+            f"V4 fallback descriptor for {kind!r} (category={category}, "
+            f"family={family}). Opaque bricks are grouped into regions "
+            "via DLPack handoff; non-opaque get inline TileLang fragments."
+        ),
+        production_source=f"cppmega_v4.models.unified_superblock_v4:BLOCK_BUILDERS[{kind!r}]",
+        production_fragment_status=(
+            "opaque_brick" if is_opaque else "v4_descriptor_synth_pending"
+        ),
+        production_fragment_reason=(
+            "auto-synthesized fallback; the brick already has a working "
+            "Path B / mlx-lm SDPA implementation that the runtime calls "
+            "via DLPack handoff. A future change can replace this with a "
+            "real TileLang fragment when the schedule template lands."
+        ),
+    )
+
+
+def build_v4_extended_registry() -> PathCBrickScheduleDescriptorRegistry:
+    """Return a registry covering both legacy v3 bricks and all V4 bricks.
+
+    Existing v3 descriptors win — we never clobber. Every key in
+    BLOCK_BUILDERS that isn't already present gets a synthesised fallback
+    descriptor. The result is a single registry that
+    ``compile_path_c_region`` can use without ``None`` lookups.
+    """
+    registry = default_path_c_brick_schedule_descriptor_registry()
+    for kind in BLOCK_BUILDERS:
+        if registry.descriptor_for(kind) is not None:
+            continue
+        registry.register(synthesize_descriptor_for_brick(kind))
+    return registry
+
+
+__all__ = [
+    "build_v4_extended_registry",
+    "synthesize_descriptor_for_brick",
+]

tests/v4/test_fusion_stage_b.py

@@ -0,0 +1,126 @@
+"""Stage B tests — descriptor_synthesizer + extended registry."""
+
+from __future__ import annotations
+
+import pytest
+
+from cppmega_mlx.runtime.path_c_fusion_schedules import (
+    PathCBrickScheduleDescriptor,
+    PathCBrickScheduleDescriptorRegistry,
+    default_path_c_brick_schedule_descriptor_registry,
+)
+from cppmega_v4.fusion.descriptor_synthesizer import (
+    build_v4_extended_registry,
+    synthesize_descriptor_for_brick,
+)
+from cppmega_v4.models.unified_superblock_v4 import BLOCK_BUILDERS
+
+
+# ---------------------------------------------------------------------------
+# Unit: per-brick synthesis
+# ---------------------------------------------------------------------------
+
+
+def test_synth_for_known_brick_returns_descriptor():
+    d = synthesize_descriptor_for_brick("gdn")
+    assert isinstance(d, PathCBrickScheduleDescriptor)
+    assert d.op_name == "gdn"
+    assert d.schedule_family == "linear_attn_scan"
+    assert d.supports_backward is True
+    assert len(d.required_codegen_steps) >= 1
+
+
+def test_synth_for_sdpa_brick_is_opaque():
+    d = synthesize_descriptor_for_brick("gated_attention")
+    assert d.schedule_family == "opaque_sdpa_with_outputs"
+    assert d.implementation_status == "opaque_brick_passthrough"
+    assert d.supports_backward is False
+
+
+def test_synth_for_unknown_kind_returns_safe_passthrough():
+    """Synthesizer must never raise for unknown kinds — produces an
+    opaque_passthrough descriptor instead."""
+    d = synthesize_descriptor_for_brick("totally_made_up_kind")
+    assert d.op_name == "totally_made_up_kind"
+    assert d.schedule_family == "opaque_passthrough"
+
+
+@pytest.mark.parametrize("kind", sorted(BLOCK_BUILDERS.keys()))
+def test_synth_for_every_block_builder_kind_is_structurally_valid(kind):
+    """Every kind in BLOCK_BUILDERS must produce a registry-acceptable
+    descriptor (non-empty op_name, codegen steps, family)."""
+    d = synthesize_descriptor_for_brick(kind)
+    assert d.op_name == kind
+    assert d.required_codegen_steps  # non-empty
+    assert d.schedule_family
+    # registry.register raises on invalid descriptors
+    PathCBrickScheduleDescriptorRegistry().register(d)
+
+
+# ---------------------------------------------------------------------------
+# Unit: build_v4_extended_registry
+# ---------------------------------------------------------------------------
+
+
+def test_extended_registry_includes_v3_legacy_descriptors():
+    reg = build_v4_extended_registry()
+    # The v3 legacy entries from default_path_c_brick_schedule_descriptor_registry
+    legacy = ("mamba3_mimo", "residual_rmsnorm", "m2rnn",
+              "attention_qkv_projection", "sparse_mla_fp8_apply")
+    for op in legacy:
+        d = reg.descriptor_for(op)
+        assert d is not None, f"legacy descriptor {op!r} dropped"
+        # legacy ones have a different status — not the v4 synth tag
+        assert d.implementation_status != "opaque_brick_passthrough"
+
+
+def test_extended_registry_does_not_clobber_legacy_when_kind_collides():
+    """If a future BLOCK_BUILDERS key matches a legacy op_name, the
+    legacy descriptor wins (descriptor_for returns it, not the synth)."""
+    reg = build_v4_extended_registry()
+    legacy_status = default_path_c_brick_schedule_descriptor_registry()\
+        .descriptor_for("residual_rmsnorm").implementation_status
+    assert reg.descriptor_for("residual_rmsnorm").implementation_status == legacy_status
+
+
+def test_extended_registry_covers_every_block_builder_kind():
+    reg = build_v4_extended_registry()
+    for kind in BLOCK_BUILDERS:
+        d = reg.descriptor_for(kind)
+        assert d is not None, f"missing descriptor for {kind!r}"
+
+
+def test_extended_registry_aot_backward_synthesis_only_for_supporting_kinds():
+    """The registry has a built-in :code:`<op>_bwd` synthesizer that only
+    fires for descriptors with supports_backward=True. Verify the v4
+    fallback descriptors respect that contract."""
+    reg = build_v4_extended_registry()
+    # gdn: supports_backward=True (linear_attn category)
+    assert reg.descriptor_for("gdn_bwd") is not None
+    # gated_attention: supports_backward=False (sdpa_attention category)
+    assert reg.descriptor_for("gated_attention_bwd") is None
+
+
+# ---------------------------------------------------------------------------
+# System: signature lookup for a real Qwen3-Next-shaped chain
+# ---------------------------------------------------------------------------
+
+
+def test_system_descriptors_for_qwen3_next_chain():
+    """Build the descriptor list for a Qwen3-Next-style 5-brick chain."""
+    reg = build_v4_extended_registry()
+    signature = ("gdn", "gdn", "gdn", "gated_attention", "moe")
+    descriptors = reg.descriptors_for_signature(signature)
+    assert descriptors is not None
+    assert len(descriptors) == len(signature)
+    for op_name, d in zip(signature, descriptors):
+        assert d.op_name == op_name
+
+
+def test_system_descriptors_for_ling26_chain():
+    """Ling 2.6's 7:1 bailing_linear + bailing_mla + bailing_moe."""
+    reg = build_v4_extended_registry()
+    signature = ("bailing_linear",) * 7 + ("bailing_mla", "bailing_moe")
+    descriptors = reg.descriptors_for_signature(signature)
+    assert descriptors is not None
+    assert len(descriptors) == 9