feat: implement backward pass support for sparse MLA FP8 path-C fusion

Author: apstenku123

cppmega_mlx/models/hybrid_lm.py

@@ -39,8 +39,9 @@
 from cppmega_mlx.nn.ngram_hash import NgramHashEmbedding
 from cppmega_mlx.nn.structure_embedding import CppMegaStructureEmbedding
 from cppmega_mlx.recipes.pattern import ExpandedNamPattern, NamLayer, expand_nam_pattern
-from cppmega_mlx.runtime.path_c_physical_abi import PathCLogicalBufferOwner
 from cppmega_mlx.runtime.kernel_policy import KernelPath, selected_path
+from cppmega_mlx.runtime.path_c_physical_abi import PathCLogicalBufferOwner
+from cppmega_mlx.runtime.path_c_taps import emit_and_tap_path_c_tensor
 from cppmega_mlx.training.mtp import MinimalMTPHead, MTPLossConfig
 
 if TYPE_CHECKING:
@@ -88,8 +89,10 @@ def __init__(
         aliases: Mapping[str, str | Sequence[str]] | None = None,
         *,
         owner_name: str = "HybridTinyLM.path_c_activation_capture",
+        capture_gradients: bool = True,
     ) -> None:
         self.owner_name = owner_name
+        self.capture_gradients = capture_gradients
         self.aliases = {
             str(source): (str(target),)
             if isinstance(target, str)
@@ -106,7 +109,13 @@ def __call__(self, event: Mapping[str, Any]) -> None:
         logical_names = tuple(str(name) for name in event.get("logical_names", ()))
         for name in logical_names:
             self.buffers[name] = tensor
-            for alias in self.aliases.get(name, ()):
+            aliases = list(self.aliases.get(name, ()))
+            if name.endswith("_grad"):
+                aliases.extend(
+                    f"{alias}_grad"
+                    for alias in self.aliases.get(name[: -len("_grad")], ())
+                )
+            for alias in aliases:
                 self.buffers[alias] = tensor
         self.events.append(event)
 
@@ -507,22 +516,23 @@ def _path_c_activation_logical_names(self, name: str) -> tuple[str, ...]:
             return (f"{brick_name}_hidden_after",)
         return (f"{brick_name}_{name}",)
 
-    def _emit_path_c_activation(self, name: str, tensor: mx.array) -> None:
+    def _emit_path_c_activation(self, name: str, tensor: mx.array) -> mx.array:
         probe = self._path_c_activation_probe()
         if probe is None:
-            return
-        probe(
-            {
+            return tensor
+        return emit_and_tap_path_c_tensor(
+            tensor,
+            probe=probe,
+            event={
                 "name": name,
                 "logical_names": self._path_c_activation_logical_names(name),
-                "tensor": tensor,
                 "layer_number": self.layer.number,
                 "layer_index": self.path_c_layer_index,
                 "route_symbol": self.layer.symbol,
                 "backend": self.backend,
                 "brick_name": self.path_c_brick_name,
                 "profile_brick_name": getattr(self, "path_c_profile_brick_name", None),
-            }
+            },
         )
 
     @property
@@ -560,17 +570,17 @@ def __call__(
     ) -> mx.array:
         self.validate_backend()
         residual = hidden_states
-        self._emit_path_c_activation("hidden", residual)
+        residual = self._emit_path_c_activation("hidden", residual)
         delta = self.route_delta(
             hidden_states,
             mask,
             kv_cache=kv_cache,
             attention_layer_idx=attention_layer_idx,
             doc_ids=doc_ids,
         )
-        self._emit_path_c_activation("delta", delta)
+        delta = self._emit_path_c_activation("delta", delta)
         updated = residual + delta
-        self._emit_path_c_activation("hidden_after", updated)
+        updated = self._emit_path_c_activation("hidden_after", updated)
         if self.mhc is not None:
             return self.mhc([updated, residual])
         return updated
@@ -628,7 +638,7 @@ def route_delta(
         if kv_cache is not None and self.backend != "attention":
             raise ValueError("kv_cache may only be passed to attention route blocks")
         x = self.norm(hidden_states)
-        self._emit_path_c_activation("normed", x)
+        x = self._emit_path_c_activation("normed", x)
         if self.backend == "attention":
             delta = cast(CausalSelfAttention, self.block)(
                 x,
@@ -637,16 +647,33 @@ def route_delta(
                 layer_idx=attention_layer_idx,
             )
         elif self.backend == "mamba3":
-            delta, _ = cast(Mamba3ReferenceBlock, self.block)(x)
+            mamba3 = cast(Mamba3ReferenceBlock, self.block)
+            probe = self._path_c_activation_probe()
+            if probe is not None:
+                h0 = mamba3.initial_h0(x.shape[0], x.dtype)
+                h0 = self._emit_path_c_activation("mamba3_h0", h0)
+                h0 = self._emit_path_c_activation("state_in", h0)
+                delta, state = mamba3(x, h0=h0)
+                self._emit_path_c_activation("state", state)
+            else:
+                delta, _ = mamba3(x)
         elif self.backend == "moe":
             delta = cast(ReferenceMoE, self.block)(x).output
         elif self.backend == "m2rnn":
             m2rnn = cast(M2RNNMixer, self.block)
-            if selected_path("m2rnn") is KernelPath.PATH_C:
-                delta, _ = m2rnn(
+            use_explicit_state = (
+                selected_path("m2rnn") is KernelPath.PATH_C
+                or self._path_c_activation_probe() is not None
+            )
+            if use_explicit_state:
+                h0 = m2rnn.initial_h0(x.shape[0], x.dtype)
+                h0 = self._emit_path_c_activation("m2rnn_h0", h0)
+                delta, state = m2rnn(
                     x,
-                    h0=m2rnn.initial_h0(x.shape[0], x.dtype),
+                    h0=h0,
+                    return_state=True,
                 )
+                self._emit_path_c_activation("m2rnn_conv_state", state.conv_state)
             else:
                 delta, _ = m2rnn(x)
         elif self.backend == "engram":
@@ -744,6 +771,20 @@ def attach_path_c_activation_probe(self, probe: PathCActivationProbe) -> int:
             raise TypeError("probe must be callable")
         for block in self.layers:
             block._path_c_activation_probe_callback = probe  # type: ignore[attr-defined]
+            attention = block.attention_block
+            if attention is not None and callable(
+                getattr(attention, "attach_path_c_prepared_probe", None)
+            ):
+                attention.attach_path_c_prepared_probe(
+                    probe,
+                    logical_prefix=str(
+                        getattr(
+                            block,
+                            "path_c_profile_brick_name",
+                            block.path_c_brick_name,
+                        )
+                    ),
+                )
         return len(self.layers)
 
     def detach_path_c_activation_probe(self) -> None:
@@ -752,6 +793,11 @@ def detach_path_c_activation_probe(self) -> None:
         for block in self.layers:
             if hasattr(block, "_path_c_activation_probe_callback"):
                 delattr(block, "_path_c_activation_probe_callback")
+            attention = block.attention_block
+            if attention is not None and callable(
+                getattr(attention, "detach_path_c_prepared_probe", None)
+            ):
+                attention.detach_path_c_prepared_probe()
 
     def path_c_parameter_logical_aliases(self) -> dict[str, tuple[str, ...]]:
         """Map MLX parameter tree names to Path C logical input names."""

cppmega_mlx/nn/_tilelang/sparse_mla_fp8_path_c.py

@@ -29,6 +29,7 @@
 
 from dataclasses import dataclass
 from functools import lru_cache
+from collections.abc import Callable, Mapping
 from typing import Any, cast
 
 import mlx.core as mx
@@ -99,6 +100,7 @@
 _SMFP8_IQKR_VEC = _SMFP8_QKR_DEFAULT_VEC
 _SMFP8_IQKR_BLOCK_K = _SMFP8_IQKR_RT * _SMFP8_IQKR_VEC
 _SMFP8_IQKR_K_WORDS = _SMFP8_IQKR_K // 4
+_SMFP8_IQKR_INDEX_DTYPE = "int32"
 
 _SMFP8_APPLY_B = 1
 _SMFP8_APPLY_S = 1
@@ -123,6 +125,8 @@
 )
 _SMFP8_APPLY_OUT_SIZE = _SMFP8_APPLY_LANES * _SMFP8_APPLY_DV
 _SMFP8_APPLY_LSE_SIZE = _SMFP8_APPLY_LANES
+_SMFP8_APPLY_OUT_DTYPE = "float32"
+_SMFP8_APPLY_INDEX_DTYPE = "int32"
 
 _SMFP8_BWD_B = 1
 _SMFP8_BWD_S = 1
@@ -143,9 +147,17 @@
 _SMFP8_BWD_IDX_SIZE = _SMFP8_BWD_B * _SMFP8_BWD_S * _SMFP8_BWD_G * _SMFP8_BWD_TOPK
 _SMFP8_BWD_DOUT_SIZE = _SMFP8_BWD_LANES * _SMFP8_BWD_DV
 _SMFP8_BWD_DOUT_DTYPE = "float32"
+_SMFP8_BWD_INDEX_DTYPE = "int32"
+_SMFP8_BWD_CLEAR_BATCH = _SMFP8_BWD_B
+_SMFP8_BWD_CLEAR_SEQ_LEN_KV = _SMFP8_BWD_SKV
+_SMFP8_BWD_CLEAR_KV_GROUP = _SMFP8_BWD_G
+_SMFP8_BWD_CLEAR_K = _SMFP8_BWD_K
 _SMFP8_BWD_CLEAR_TOTAL = _SMFP8_BWD_B * _SMFP8_BWD_SKV * _SMFP8_BWD_G * _SMFP8_BWD_K
 _SMFP8_BWD_CLEAR_THREADS = 256
 _SMFP8_PER_TOKEN_QUANT_THREADS = 256
+_SMFP8_PTQ_ROWS = 1
+_SMFP8_PTQ_K = 64
+_SMFP8_PTQ_INPUT_DTYPE = "float32"
 _SMFP8_PREPARE_Q_ROWS = 1
 _SMFP8_PREPARE_KV_ROWS = 1
 _SMFP8_PREPARE_ROWS = _SMFP8_PREPARE_Q_ROWS + _SMFP8_PREPARE_KV_ROWS
@@ -208,6 +220,24 @@ class SparseMLAFp8PathCDirectError(RuntimeError):
     """Raised when a prepared-buffer tvm-ffi owner-output path cannot run."""
 
 
+def _emit_fp8_apply_runtime_buffer(
+    probe: Callable[[Mapping[str, Any]], None] | None,
+    *,
+    name: str,
+    tensor: mx.array,
+) -> None:
+    if probe is None:
+        return
+    probe(
+        {
+            "name": name,
+            "tensor": tensor,
+            "producer_owner": "sparse_mla_fp8_path_c_apply",
+            "producer_stage": "sparse_mla_fp8_apply",
+        }
+    )
+
+
 def _index_dtype_name(indices: mx.array, *, op_name: str) -> str:
     if indices.dtype == mx.int32:
         return "int32"
@@ -1848,9 +1878,6 @@ def fp8_sparse_mla_apply_kernel(
             inv_sum = T.alloc_local((1,), "float32")
             stride = T.alloc_local((1,), "int32")
             gather_idx = T.alloc_local((1,), "int32")
-            kv_row_base_local = T.alloc_local((1,), "int32")
-            kv_scale_idx_local = T.alloc_local((1,), "int32")
-            dkv_idx_local = T.alloc_local((1,), "int32")
 
             h = bx % _SMFP8_APPLY_H
             b = bx // (_SMFP8_APPLY_H * _SMFP8_APPLY_S)
@@ -2984,6 +3011,7 @@ def sparse_mla_fp8_path_c_apply(
     out: mx.array | None = None,
     lse: mx.array | None = None,
     output_dtype: mx.Dtype | None = None,
+    runtime_buffer_probe: Callable[[Mapping[str, Any]], None] | None = None,
 ) -> mx.array | tuple[mx.array, mx.array] | None:
     """Run fused FP8 Sparse-MLA Path C over prepared GPU buffers.
 
@@ -3071,6 +3099,21 @@ def sparse_mla_fp8_path_c_apply(
                 raise ValueError("sinks must be float32")
             sinks_buf = sinks
             has_sinks_buf = mx.array([1], dtype=mx.int32)
+        _emit_fp8_apply_runtime_buffer(
+            runtime_buffer_probe,
+            name="sparse_mla_sm_scale",
+            tensor=sm_scale_buf,
+        )
+        _emit_fp8_apply_runtime_buffer(
+            runtime_buffer_probe,
+            name="sparse_mla_sinks",
+            tensor=sinks_buf,
+        )
+        _emit_fp8_apply_runtime_buffer(
+            runtime_buffer_probe,
+            name="sparse_mla_has_sinks",
+            tensor=has_sinks_buf,
+        )
         returned = kernel(
             _flat_1d_view(q_fp8),
             _flat_1d_view(q_scale),
@@ -3093,8 +3136,13 @@ def sparse_mla_fp8_path_c_apply(
             raise RuntimeError(
                 "sparse_mla_fp8_path_c_apply: native TileLang tvm-ffi "
                 "graph-output dispatch did not return out/lse"
-            )
+        )
         return None
+    _emit_fp8_apply_runtime_buffer(
+        runtime_buffer_probe,
+        name="lse",
+        tensor=cast(mx.array, returned[1]),
+    )
     direct_out = cast(mx.array, returned[0]).reshape(
         (batch, seq_len, heads, d_v_resolved)
     )
@@ -3250,6 +3298,7 @@ def sparse_mla_fp8_path_c_apply_prepared_float(
     force_path_c: bool = False,
     causal: bool = False,
     output_dtype: mx.Dtype | None = None,
+    runtime_buffer_probe: Callable[[Mapping[str, Any]], None] | None = None,
 ) -> mx.array:
     """Differentiable owner wrapper for prepared-buffer FP8 Path C apply.
 
@@ -3275,6 +3324,7 @@ def _apply(
             sinks=sinks,
             force_path_c=force_path_c,
             output_dtype=_mx_float_dtype(output_dtype, default=q_in.dtype),
+            runtime_buffer_probe=runtime_buffer_probe,
         )
         if out is None:
             if force_path_c: