feat(v4): KDA Paths B/C/D — mirror GDN's real implementations

Three new modules under cppmega_v4/_tilelang/:

  - kda_path_b.py — hand-MSL recurrent forward via mx.fast.metal_kernel.
    Per-thread layout: one thread per (b, hv_idx, v_index). Each thread
    owns one column of the K x V state, decays per-K via exp(g[B,T,HV,K])
    (vectorized gate, KDA-specific), reduces over K for KS dot, computes
    β-corrected delta, rank-1 outer adds, then projects q · S. Supports
    arbitrary (B, T, H, HV, K, V) with HV % H == 0. Falls back to Path A
    for initial_state, custom scale, or HV not divisible by H.

  - kda_path_c.py — TileLang DSL @T.prim_func lowered through
    cppmega_mlx's dispatch_lower + tilelang.compile(target='metal',
    execution_backend='tvm_ffi'). Per-lane recurrent scan modeled on
    mamba3_path_c.py: LANES = B * HV * V threads, state held in
    registers as K floats. Host TileLang→MSL infra imported read-only.

  - kda_path_d.py — Triton-frontend seam over FLA KDA chunk kernels
    (11 prims, ~3340 LoC). Two-stage probe (triton+frontend, then
    fla.ops.kda.chunk); status names the actionable blocker
    (op_mapping needs matmul/exp/masked-load emitters).

kda_paths.py rewritten to wire B/C/D into real dispatch (previously
all three fell through to Path A). Auto-mode preference unchanged
(C > B > D > A).

Tests: 16 new in tests/v4/test_kda_paths.py covering all three paths:
  - B: status reports Metal kernel; parity vs Path A (atol 1e-4);
    final-state return; env-forced dispatch.
  - C: module import w/o tilelang; status text precision;
    runtime↔dispatch round-trip; env-forced dispatch returns valid
    output; fallback parity bit-exact when tilelang missing.
  - D: module import; concrete blocker named; runtime↔dispatch
    round-trip; both probes return tuples; env-forced fallback;
    lowering seam returns (None, msg).
  + 1 dispatch sanity test (path keys unchanged).

Existing test_kda_dispatch_returns_same_as_path_a relaxed to
assert_allclose(atol=1e-5) since auto-mode now picks Path B (real
Metal float32) over Path A (MLX float64-then-cast).

v4 suite: 153 passed / 11 skipped.

Author: apstenku123

cppmega_v4/_tilelang/kda_path_b.py

@@ -0,0 +1,194 @@
+"""KDA Path B — hand-MSL recurrent forward via mx.fast.metal_kernel.
+
+Mirrors ``linear_attention_path_b.py`` but for the KDA recurrence (FLA naive):
+
+    q, k: [B, T, H, K]         (q scaled by 1/sqrt(K), then both repeat to HV)
+    v:    [B, T, HV, V]
+    g:    [B, T, HV, K]        (per-K vectorized log-gate)
+    beta: [B, T, HV]
+    S:    [B, HV, K, V]
+
+    for t in [0, T):
+        S       *= exp(g_t)               # per-K decay, broadcast over V
+        inner    = v_t - sum_k(k_t * S)
+        S       += (beta_t * k_t)[:, :, None] * inner[:, None, :]
+        o_t      = sum_k(q_t * S)
+
+Per-thread layout: one thread per (batch, hv_head, v_index). Each thread
+owns one column of the K x V state for one (b, hv) and walks time
+serially while reducing over K for state-derived quantities.
+
+Group expansion (HV / H): each ``hv`` maps to ``h = hv // G`` for indexing
+into q and k. The kernel performs that mapping inline.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+from cppmega_v4._tilelang._kernel_cache import get_or_build_kernel
+
+
+def _kda_forward_kernel(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+) -> tuple[mx.array, mx.array]:
+    """Metal forward for KDA recurrence.
+
+    Args:
+        q, k: [B, T, H, K]
+        v:    [B, T, HV, V]
+        g:    [B, T, HV, K]   per-K vectorized log-gate
+        beta: [B, T, HV]
+
+    Returns:
+        o: [B, T, HV, V]
+        S_final: [B, HV, K, V]
+    """
+    if q.ndim != 4 or k.shape != q.shape:
+        raise ValueError(f"q/k must be [B,T,H,K]; got q={q.shape}, k={k.shape}")
+    if v.ndim != 4 or v.shape[:2] != q.shape[:2]:
+        raise ValueError(f"v must be [B,T,HV,V]; got v={v.shape}")
+    if g.shape != (*v.shape[:3], k.shape[-1]):
+        raise ValueError(f"g must be [B,T,HV,K]; got g={g.shape}")
+    if beta.shape != v.shape[:3]:
+        raise ValueError(f"beta must be [B,T,HV]; got beta={beta.shape}")
+
+    b, t, h, kdim = q.shape
+    hv, vdim = v.shape[2], v.shape[-1]
+    if hv % h != 0:
+        raise ValueError(f"HV ({hv}) must be divisible by H ({h})")
+    group = hv // h
+
+    q = q.astype(mx.float32) * (kdim ** -0.5)
+    k = k.astype(mx.float32)
+    v = v.astype(mx.float32)
+    g = g.astype(mx.float32)
+    beta = beta.astype(mx.float32)
+
+    q_flat = q.reshape(-1)
+    k_flat = k.reshape(-1)
+    v_flat = v.reshape(-1)
+    g_flat = g.reshape(-1)
+    beta_flat = beta.reshape(-1)
+
+    source = f"""
+        uint vj  = thread_position_in_grid.x;
+        uint bhv = thread_position_in_grid.y;
+
+        if (vj >= {vdim}u || bhv >= {b * hv}u) return;
+
+        uint bb     = bhv / {hv}u;
+        uint hv_idx = bhv % {hv}u;
+        uint h_idx  = hv_idx / {group}u;
+
+        // Per-thread state column: S[bb, hv_idx, :, vj]  size K
+        float state[{kdim}];
+        for (int i = 0; i < {kdim}; i++) state[i] = 0.0f;
+
+        for (int ti = 0; ti < {t}; ti++) {{
+            // g_base: g[bb, ti, hv_idx, 0]
+            int g_base    = ((bb * {t} + ti) * {hv} + hv_idx) * {kdim};
+            // beta scalar: beta[bb, ti, hv_idx]
+            int beta_idx  = (bb * {t} + ti) * {hv} + hv_idx;
+            float beta_t  = beta[beta_idx];
+
+            // k/q base for this (b, ti, h_idx)
+            int qk_base   = ((bb * {t} + ti) * {h} + h_idx) * {kdim};
+            // v scalar: v[bb, ti, hv_idx, vj]
+            int v_idx     = ((bb * {t} + ti) * {hv} + hv_idx) * {vdim} + vj;
+            float v_j     = v[v_idx];
+
+            // Phase 1: per-K decay  AND  KS reduction along K (interleaved)
+            float kth_S_j = 0.0f;
+            for (int i = 0; i < {kdim}; i++) {{
+                float decay_i = exp(g[g_base + i]);
+                state[i] *= decay_i;
+                kth_S_j += k[qk_base + i] * state[i];
+            }}
+
+            // Phase 2: delta correction
+            float inner_j = v_j - kth_S_j;
+
+            // Phase 3: rank-1 outer add  S[i,j] += beta * k[i] * inner  AND  o[j] = sum_i q[i] * S[i,j]
+            float o_j = 0.0f;
+            for (int i = 0; i < {kdim}; i++) {{
+                float k_i = k[qk_base + i];
+                float q_i = q[qk_base + i];
+                state[i] += beta_t * k_i * inner_j;
+                o_j += q_i * state[i];
+            }}
+
+            // output[bb, ti, hv_idx, vj]
+            int o_idx = ((bb * {t} + ti) * {hv} + hv_idx) * {vdim} + vj;
+            output[o_idx] = o_j;
+        }}
+
+        // Final state column: S_final[bb, hv_idx, :, vj]
+        int sf_base = (bb * {hv} + hv_idx) * {kdim * vdim} + vj;
+        for (int i = 0; i < {kdim}; i++) {{
+            state_final[sf_base + i * {vdim}] = state[i];
+        }}
+    """
+
+    name = f"v4_kda_fwd_{b}_{t}_{h}_{hv}_{kdim}_{vdim}"
+    kernel = get_or_build_kernel(
+        name=name,
+        input_names=["q", "k", "v", "g", "beta"],
+        output_names=["output", "state_final"],
+        source=source,
+    )
+
+    grid = (vdim, b * hv, 1)
+    tg_x = min(vdim, 64)
+    threadgroup = (tg_x, 1, 1)
+
+    out, sf = kernel(
+        inputs=[q_flat, k_flat, v_flat, g_flat, beta_flat],
+        output_shapes=[
+            (b * t * hv * vdim,),
+            (b * hv * kdim * vdim,),
+        ],
+        output_dtypes=[mx.float32, mx.float32],
+        grid=grid,
+        threadgroup=threadgroup,
+    )
+    return out.reshape(b, t, hv, vdim), sf.reshape(b, hv, kdim, vdim)
+
+
+def kda_forward_path_b(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    *,
+    scale: float | None = None,
+    initial_state: mx.array | None = None,
+    output_final_state: bool = False,
+):
+    """KDA Path B forward, signature matching ``naive_recurrent_kda``.
+
+    Falls back to Path A for unsupported configs (initial_state, custom
+    scale, or HV not divisible by H) so the dispatch never produces wrong
+    numerics silently.
+    """
+    if (
+        initial_state is not None
+        or scale is not None
+        or v.shape[2] % q.shape[2] != 0
+    ):
+        from cppmega_v4.nn._external.fla_naive_kda import naive_recurrent_kda
+        return naive_recurrent_kda(
+            q, k, v, g, beta,
+            scale=scale, initial_state=initial_state,
+            output_final_state=output_final_state,
+        )
+    o, sf = _kda_forward_kernel(q, k, v, g, beta)
+    return o, (sf if output_final_state else None)
+
+
+__all__ = ["kda_forward_path_b"]

cppmega_v4/_tilelang/kda_path_c.py

@@ -0,0 +1,175 @@
+"""KDA Path C — TileLang DSL ``@T.prim_func`` lowered to Metal via tvm_ffi.
+
+Mirrors ``linear_attention_path_c.py`` but for the KDA recurrence:
+
+    q, k: [B, T, H, K]    (q pre-scaled by 1/sqrt(K), then both repeat to HV)
+    v:    [B, T, HV, V]
+    g:    [B, T, HV, K]   per-K vectorized log-gate
+    beta: [B, T, HV]
+    S:    [B, HV, K, V]
+
+Per-lane scan: one lane per (b, hv, v_idx). State held in registers as K
+floats per lane. Inner loop: per-K decay + KS reduction (interleaved),
+δ correction, rank-1 outer add + q·S projection.
+
+Plugin invariant: imports ``dispatch_lower`` / ``_msl_transform`` from the
+host ``cppmega_mlx`` package read-only — never modifies them.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Any
+
+import mlx.core as mx
+
+
+def _tilelang_importable() -> tuple[bool, str]:
+    try:
+        import tilelang  # noqa: F401
+        import tilelang.language as _T  # noqa: F401
+        from tilelang.engine.lower import lower as _lower  # noqa: F401
+    except Exception as exc:
+        return False, f"tilelang import failed: {exc.__class__.__name__}: {exc}"
+    try:
+        from cppmega_mlx.nn._tilelang._engine_dispatch import dispatch_lower  # noqa: F401
+        from cppmega_mlx.nn._tilelang import _msl_transform  # noqa: F401
+    except Exception as exc:
+        return False, f"host TileLang→MSL infra not reachable: {exc}"
+    return True, "tilelang + host TileLang→MSL infra reachable"
+
+
+def _threads_for(lanes: int) -> int:
+    for tg in (256, 192, 128, 96, 64, 32):
+        if lanes % tg == 0:
+            return tg
+    return 32
+
+
+@lru_cache(maxsize=64)
+def _kda_fwd_kernel_for(
+    BATCH: int,
+    SEQ: int,
+    HEADS: int,
+    HV: int,
+    HEADDIM_K: int,
+    HEADDIM_V: int,
+    q_dtype: str = "float32",
+    k_dtype: str = "float32",
+    v_dtype: str = "float32",
+    g_dtype: str = "float32",
+    beta_dtype: str = "float32",
+    h0_dtype: str = "float32",
+    y_dtype: str = "float32",
+    h_last_dtype: str = "float32",
+) -> tuple[Any, Any]:
+    """Build (and cache) the KDA fwd Path C kernel for this shape/dtype tuple."""
+    import tilelang
+    import tilelang.language as T
+    from cppmega_mlx.nn._tilelang import _msl_transform
+    from cppmega_mlx.nn._tilelang._engine_dispatch import dispatch_lower
+
+    assert HV % HEADS == 0, f"HV ({HV}) must be divisible by HEADS ({HEADS})"
+    GROUP = HV // HEADS
+    LANES = BATCH * HV * HEADDIM_V
+    THREADS = _threads_for(LANES)
+    accum_dtype = "float32"
+
+    @T.prim_func
+    def fwd(
+        q: T.Tensor((BATCH, SEQ, HEADS, HEADDIM_K), q_dtype),
+        k: T.Tensor((BATCH, SEQ, HEADS, HEADDIM_K), k_dtype),
+        v: T.Tensor((BATCH, SEQ, HV, HEADDIM_V), v_dtype),
+        g: T.Tensor((BATCH, SEQ, HV, HEADDIM_K), g_dtype),
+        beta: T.Tensor((BATCH, SEQ, HV), beta_dtype),
+        h0: T.Tensor((BATCH, HV, HEADDIM_K, HEADDIM_V), h0_dtype),
+        y: T.Tensor((BATCH, SEQ, HV, HEADDIM_V), y_dtype),
+        h_last: T.Tensor((BATCH, HV, HEADDIM_K, HEADDIM_V), h_last_dtype),
+    ):
+        with T.Kernel(T.ceildiv(LANES, THREADS), threads=THREADS) as bx:
+            tid = T.get_thread_binding(0)
+            global_lane = bx * THREADS + tid
+            h_state = T.alloc_local((HEADDIM_K,), accum_dtype)
+            if global_lane < LANES:
+                vj = global_lane % HEADDIM_V
+                hv_idx = (global_lane // HEADDIM_V) % HV
+                bb = global_lane // (HEADDIM_V * HV)
+                h_idx = hv_idx // GROUP
+                for i in T.serial(HEADDIM_K):
+                    h_state[i] = T.cast(h0[bb, hv_idx, i, vj], accum_dtype)
+                for t in T.serial(SEQ):
+                    beta_val = T.cast(beta[bb, t, hv_idx], accum_dtype)
+                    v_j = T.cast(v[bb, t, hv_idx, vj], accum_dtype)
+                    # Phase 1: per-K decay + KS reduction (interleaved).
+                    kth_S_j = T.alloc_var(T.float32, init=0.0)
+                    for i in T.serial(HEADDIM_K):
+                        decay_i = T.exp(T.cast(g[bb, t, hv_idx, i], accum_dtype))
+                        h_state[i] = h_state[i] * decay_i
+                        kth_S_j += T.cast(k[bb, t, h_idx, i], accum_dtype) * h_state[i]
+                    inner_j = v_j - kth_S_j
+                    # Phase 2: rank-1 outer add + q·h projection.
+                    out = T.alloc_var(T.float32, init=0.0)
+                    for i in T.serial(HEADDIM_K):
+                        k_i = T.cast(k[bb, t, h_idx, i], accum_dtype)
+                        q_i = T.cast(q[bb, t, h_idx, i], accum_dtype)
+                        h_state[i] = h_state[i] + beta_val * k_i * inner_j
+                        out += q_i * h_state[i]
+                    y[bb, t, hv_idx, vj] = T.cast(out, y_dtype)
+                for i in T.serial(HEADDIM_K):
+                    h_last[bb, hv_idx, i, vj] = T.cast(h_state[i], h_last_dtype)
+
+    artifact = dispatch_lower(fwd, target="metal", return_msl=True)
+    kernel = tilelang.compile(
+        fwd,
+        target=_msl_transform._as_metal_target("metal"),
+        execution_backend="tvm_ffi",
+        out_idx=[6, 7],  # y, h_last
+    )
+    return kernel, artifact
+
+
+def _path_c_runtime_status() -> tuple[bool, str]:
+    return _tilelang_importable()
+
+
+def _kda_fwd_path_c_call(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    *,
+    scale: float | None = None,
+    initial_state: mx.array | None = None,
+    output_final_state: bool = False,
+):
+    """KDA Path C entry — same signature as ``naive_recurrent_kda``."""
+    import math
+
+    B, T_, H, K_dim = q.shape
+    HV, V_dim = v.shape[2], v.shape[-1]
+    fla_scale = scale if scale is not None else 1.0 / math.sqrt(K_dim)
+    q_scaled = (q.astype(mx.float32) * fla_scale).astype(q.dtype)
+
+    h0 = initial_state
+    if h0 is None:
+        h0 = mx.zeros((B, HV, K_dim, V_dim), dtype=mx.float32)
+
+    kernel, _ = _kda_fwd_kernel_for(
+        B, T_, H, HV, K_dim, V_dim,
+        q_dtype=str(q.dtype).rsplit(".", 1)[-1],
+        k_dtype=str(k.dtype).rsplit(".", 1)[-1],
+        v_dtype=str(v.dtype).rsplit(".", 1)[-1],
+        g_dtype=str(g.dtype).rsplit(".", 1)[-1],
+        beta_dtype=str(beta.dtype).rsplit(".", 1)[-1],
+    )
+    y, h_last = kernel(q_scaled, k, v, g, beta, h0)
+    return y, (h_last if output_final_state else None)
+
+
+__all__ = [
+    "_kda_fwd_kernel_for",
+    "_kda_fwd_path_c_call",
+    "_path_c_runtime_status",
+    "_tilelang_importable",
+]