feat: add comprehensive training profiling reports and updates to 1B model benchmark scripts and core NN components

Author: apstenku123

cppmega_mlx/models/hybrid_lm.py

@@ -470,6 +470,7 @@ def __call__(
         *,
         kv_cache: ContiguousKVCache | None = None,
         attention_layer_idx: int | None = None,
+        doc_ids: mx.array | None = None,
     ) -> mx.array:
         self.validate_backend()
         residual = hidden_states
@@ -478,6 +479,7 @@ def __call__(
             mask,
             kv_cache=kv_cache,
             attention_layer_idx=attention_layer_idx,
+            doc_ids=doc_ids,
         )
         updated = residual + delta
         if self.mhc is not None: