fix(v4-path-d): correct kkt BT/BC contract + with_state BV tiling (3/8 cells green; ffi-NULL unblocked)

The ffi-NULL fix (tilelang 168877b0) took Path D from 0/8 launched -> 8/8 launched, exposing two
cppmega-side launch-config bugs (RULE #1 surfaced them; no fallback):
- kkt_solve was launched BT=32,BC=16 but the FLA kkt kernel hard-codes a 4-sub-block decomposition
  requiring BT==4*BC -> wrong/NaN. Fixed BC=GDN_RUNTIME_BT//4=8 + pinned _NUM_WARPS=1 for the [8,8]
  GEMM (plumbed _NUM_WARPS through the lowering driver into from_ttir(num_warps=...)).
- with_state initial-state chunk_h exceeded Apple's compute-function stack -> lowered
  GDN_CHUNK_H_METAL_SAFE_BV 16->8 (the predicted recompute_w_u-style BV tiling).

Result: 3/8 cells green in isolation (fixed_k64_v32 0.018, custom_scale 0.009, with_state 0.019 --
named residual solved; kkt fix dropped fixed_k64 0.137->0.018). No fallback added; unsupported
shapes still RAISE PathDRuntimeUnavailable.

REMAINING (tilelang-side, separate fix): cells with chunks/seq>=3 or multi-head HV>=2 still NaN/wrong
due to the TVM-FFI Metal bridge lazy-multi-launch buffer/producer-registry id()-reuse hazard (same
metal_graph_sync issue) -- needs a tilelang bridge fix, not cppmega.

Author: jeweldave

cppmega_v4/_tilelang/linear_attention_path_d_real.py

@@ -876,6 +876,18 @@ def _lower_fla_kernel_uncached(
 
     # Prefer ``from_ttir`` directly so we control the text-vs-mlir routing;
     # ``from_triton_kernel`` would re-capture and double the work.
+    #
+    # A stage may pin the cooperative-tensor warp count via a private
+    # ``_NUM_WARPS`` constexpr (filtered from the TTIR signature above because
+    # it is ``_``-prefixed). This is required for small-tile GEMMs such as the
+    # GDN kkt_solve at BC=8: a [8, 8] block cannot be partitioned across the
+    # lowering's default 4 warps (``m_warp * n_warp must equal num_warps``), so
+    # the stage requests ``num_warps=1``. It is a launch-config override, not a
+    # change to the lowered kernel body.
+    num_warps_override = constexprs.get("_NUM_WARPS")
+    num_warps_override = (
+        int(num_warps_override) if num_warps_override is not None else None
+    )
     try:
         from poc.triton_frontend import from_ttir, _walk_text_ttir
 
@@ -887,6 +899,7 @@ def _lower_fla_kernel_uncached(
                 name=primfunc_name,
                 grid=grid,
                 arg_buffer_shapes=arg_buffer_shapes,
+                num_warps=num_warps_override,
             )
             return LowerResult(
                 status="LOWERED_FULL",

cppmega_v4/_tilelang/path_d_runtime_adapter.py

@@ -35,7 +35,17 @@
 GDN_FIXED_BT = 64
 GDN_FIXED_BV = 32
 GDN_RUNTIME_BT = 32
-GDN_CHUNK_H_METAL_SAFE_BV = 16
+# chunk_delta_h (h_blockdim64) holds [64, BV] fp32 state accumulators (b_h1..)
+# thread-private. With an initial state loaded (USE_INITIAL_STATE, the
+# ``with_state`` shape) the extra h0 load + the BV-wide store path pushes the
+# per-thread footprint past Apple's compute-function stack budget at BV=16,
+# so newComputePipelineStateWithFunction fails ("Compute function exceeds
+# available stack space"). Tiling the V reduction to BV=8 (the smallest tile
+# Metal's GEMM allows -- N must be a multiple of 8) shrinks the state tile so
+# the pipeline-state creation fits, mirroring the recompute_w_u BK/BV tiling
+# below. BV=8 still launches and stays in parity for the no-initial-state
+# shapes, so it is used uniformly.
+GDN_CHUNK_H_METAL_SAFE_BV = 8
 GDN_FIXED_CHUNK_O_BV = 16
 # recompute_w_u (shared by GDN and KDA) materializes several BT*BK + BT*BV
 # THREAD-PRIVATE accumulators (b_w, b_u, b_kb, b_vb, b_A, ...). With the FLA
@@ -1926,7 +1936,27 @@ def _compile_gdn_runtime_stages(
         "BT": GDN_RUNTIME_BT,
         "IS_VARLEN": bool(is_varlen),
     }
-    kkt_constexprs = {**runtime_private, **shared_dims}
+    # The FLA kkt_solve kernel hard-codes a 4-sub-block (BC x 4) decomposition
+    # of each BT-token chunk: it reads/inverts/stores blocks at column offsets
+    # 0, BC, 2*BC, 3*BC, so it is only correct when BT == 4 * BC. The default
+    # FLA pairing is BT=64 / BC=16. When this multi-stage launcher runs the
+    # whole GDN pipeline at the Metal-safe GDN_RUNTIME_BT (=32, because the
+    # chunk_o / chunk_h inter-chunk kernels overflow Apple's 32 KiB threadgroup
+    # memory at BT=64), the kkt stage MUST shrink BC to BT/4 so its 4 sub-blocks
+    # still tile exactly one BT chunk -- otherwise sub-blocks 2/3 read the next
+    # chunk's tokens and the (I+A)^{-1} merge writes past the [T, BT] A buffer,
+    # producing a wrong inverse that drifts across chunks (and NaNs past 2
+    # chunks). A [BC, BC] = [8, 8] GEMM cannot be split across the lowering's
+    # default 4 warps (m_warp * n_warp must equal num_warps), so the kkt stage
+    # also pins num_warps=1 for the small-BC tile. Both are launch-config
+    # corrections on the cppmega side; the lowered kernel itself is unchanged.
+    kkt_bc = max(1, GDN_RUNTIME_BT // 4)
+    kkt_constexprs = {
+        **runtime_private,
+        **shared_dims,
+        "BC": kkt_bc,
+        "_NUM_WARPS": 1 if kkt_bc < 16 else 4,
+    }
     recompute_constexprs = {
         **runtime_private,
         **shared_dims,