refactor: unify Path-C scratch buffer management with coalesced ABI banks and optimized TIR emission

Author: apstenku123

cppmega_mlx/_mlx_lm_imports.py

@@ -10,8 +10,13 @@
 
 from collections.abc import Iterator
 from contextlib import contextmanager
+from functools import lru_cache
+import sys
+from typing import Any
 import warnings
 
+import mlx.core as mx
+
 
 _SENTENCEPIECE_SWIG_WARNING_MESSAGES = (
     r"builtin type SwigPyPacked has no __module__ attribute",
@@ -39,9 +44,72 @@ def suppress_sentencepiece_swig_warnings() -> Iterator[None]:
 
 
 _install_sentencepiece_swig_filters()
-with suppress_sentencepiece_swig_warnings():
-    from mlx_lm.models.base import scaled_dot_product_attention
-    from mlx_lm.models.cache import KVCache, QuantizedKVCache
+
+
+def _load_mlx_lm_scaled_dot_product_attention() -> object:
+    if "tvm" in sys.modules and "mlx_lm.models.base" not in sys.modules:
+        raise RuntimeError(
+            "MLX-LM quantized attention cannot be imported after TVM/LLVM in "
+            "the same process; run inference and TileLang compile in separate "
+            "processes or load MLX-LM before TVM"
+        )
+    with suppress_sentencepiece_swig_warnings():
+        from mlx_lm.models.base import scaled_dot_product_attention
+
+    return scaled_dot_product_attention
+
+
+def scaled_dot_product_attention(
+    queries: mx.array,
+    keys: mx.array | tuple[mx.array, ...],
+    values: mx.array | tuple[mx.array, ...],
+    cache: Any,
+    scale: float,
+    mask: mx.array | None,
+    sinks: mx.array | None = None,
+) -> mx.array:
+    """MLX-LM-compatible SDPA without importing its Torch/Triton stack by default."""
+
+    needs_mlx_lm_quantized_path = hasattr(cache, "bits") or (
+        isinstance(keys, tuple)
+        and cache is not None
+        and cache.__class__.__name__ == "TurboQuantKVCache"
+    )
+    if needs_mlx_lm_quantized_path:
+        sdpa = _load_mlx_lm_scaled_dot_product_attention()
+        return sdpa(
+            queries,
+            keys,
+            values,
+            cache=cache,
+            scale=scale,
+            mask=mask,
+            sinks=sinks,
+        )
+    return mx.fast.scaled_dot_product_attention(
+        queries,
+        keys,
+        values,
+        scale=scale,
+        mask=mask,
+        sinks=sinks,
+    )
+
+
+@lru_cache(maxsize=1)
+def _load_kv_cache_classes() -> tuple[type[object], type[object]]:
+    with suppress_sentencepiece_swig_warnings():
+        from mlx_lm.models.cache import KVCache, QuantizedKVCache
+
+    return KVCache, QuantizedKVCache
+
+
+def __getattr__(name: str) -> object:
+    if name == "KVCache":
+        return _load_kv_cache_classes()[0]
+    if name == "QuantizedKVCache":
+        return _load_kv_cache_classes()[1]
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 
 __all__ = [

cppmega_mlx/inference/engine.py

@@ -3,14 +3,16 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, replace
-from typing import TypeAlias, cast
+from typing import TYPE_CHECKING, Any, TypeAlias, cast
 
 import mlx.core as mx
 
-from cppmega_mlx._mlx_lm_imports import KVCache, QuantizedKVCache
 from cppmega_mlx.inference.quantization import make_quantized_kv_cache
 
-_LayerCache: TypeAlias = KVCache | QuantizedKVCache
+if TYPE_CHECKING:
+    from mlx_lm.models.cache import KVCache
+
+_LayerCache: TypeAlias = Any
 _StateTree: TypeAlias = mx.array | tuple["_StateTree", ...]
 
 
@@ -53,6 +55,8 @@ class ContiguousKVCache:
     """A thin validated wrapper around one MLX-LM KV cache per layer."""
 
     def __init__(self, config: ContiguousKVCacheConfig) -> None:
+        from cppmega_mlx._mlx_lm_imports import KVCache
+
         self.config = config
         self.layers: list[_LayerCache] = [
             _make_layer_cache(config) for _ in range(config.num_layers)
@@ -311,6 +315,8 @@ def prefill_contiguous_kv_cache(
             destination.config.batch_size,
         )
         dst_layer.state = copied_state
+        from cppmega_mlx._mlx_lm_imports import QuantizedKVCache
+
         if isinstance(dst_layer, QuantizedKVCache):
             dst_layer.meta_state = src_layer.meta_state
 
@@ -340,6 +346,8 @@ def _make_layer_cache(config: ContiguousKVCacheConfig) -> _LayerCache:
         return make_quantized_kv_cache(
             bits=config.kv_bits, group_size=config.kv_group_size
         )
+    from cppmega_mlx._mlx_lm_imports import KVCache
+
     return KVCache()
 
 

cppmega_mlx/inference/quantization.py

@@ -4,10 +4,13 @@
 
 from collections.abc import Callable, Iterable
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import mlx.nn as nn
 
-from cppmega_mlx._mlx_lm_imports import KVCache, QuantizedKVCache
+if TYPE_CHECKING:
+    from mlx_lm.models.cache import KVCache, QuantizedKVCache
+
 _SUPPORTED_BITS = frozenset({4, 8})
 _SUPPORTED_GROUP_SIZES = frozenset({32, 64, 128})
 _SUPPORTED_MODES = frozenset({"affine"})
@@ -127,6 +130,8 @@ def make_quantized_kv_cache(
     """Create an mlx-lm ``QuantizedKVCache`` with q4 defaults."""
 
     _validate_quant_args(bits=bits, group_size=group_size, mode="affine")
+    from cppmega_mlx._mlx_lm_imports import QuantizedKVCache
+
     return QuantizedKVCache(group_size=group_size, bits=bits)
 
 
@@ -139,6 +144,8 @@ def quantize_kv_cache(
     """Convert an existing mlx-lm ``KVCache`` to ``QuantizedKVCache``."""
 
     _validate_quant_args(bits=bits, group_size=group_size, mode="affine")
+    from cppmega_mlx._mlx_lm_imports import KVCache, QuantizedKVCache
+
     if isinstance(cache, QuantizedKVCache):
         if cache.bits != bits or cache.group_size != group_size:
             raise ValueError(

cppmega_mlx/nn/attention.py

@@ -10,7 +10,6 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from cppmega_mlx._mlx_lm_imports import scaled_dot_product_attention
 from cppmega_mlx.inference.engine import ContiguousKVCache
 from cppmega_mlx.runtime.kernel_policy import KernelPath, record_dispatch, selected_path
 from cppmega_mlx.runtime.path_c_taps import emit_and_tap_path_c_tensor
@@ -1056,6 +1055,8 @@ def __call__(
                 key_length=key_length,
             )
         sinks = _validate_attention_sinks(sinks, self.config.num_q_heads)
+        from cppmega_mlx._mlx_lm_imports import scaled_dot_product_attention
+
         out = scaled_dot_product_attention(
             q,
             k,

cppmega_mlx/runtime/path_c_fusion.py

@@ -2916,7 +2916,7 @@ def _workspace_edge_buffers_for_attested_template(
         return ()
     if not _region_has_attention_kv_workspace_edges(region):
         return ()
-    return ("kv_fp8", "kv_scale")
+    return ("q_fp8", "q_scale", "kv_fp8", "kv_scale")
 
 
 def _region_has_attention_kv_workspace_edges(region: PathCFusionRegion) -> bool:
@@ -2925,7 +2925,8 @@ def _region_has_attention_kv_workspace_edges(region: PathCFusionRegion) -> bool:
         node_by_name[edge.producer].op_name == "attention_qkv_projection"
         and node_by_name[edge.consumer].op_name
         in {"sparse_mla_fp8_apply", "sparse_mla_fp8_apply_bwd"}
-        and _canonical_path_c_edge_buffer_name(edge.input) in {"kv_fp8", "kv_scale"}
+        and _canonical_path_c_edge_buffer_name(edge.input)
+        in {"q_fp8", "q_scale", "kv_fp8", "kv_scale"}
         for edge in region.edges
     )
 
@@ -3540,6 +3541,10 @@ def _infer_edges(nodes: Sequence[FusionNode]) -> tuple[FusionEdge, ...]:
 
 def _canonical_path_c_edge_buffer_name(buffer_name: str) -> str:
     name = str(buffer_name)
+    if name.endswith("_q_fp8"):
+        return "q_fp8"
+    if name.endswith("_q_scale"):
+        return "q_scale"
     if name.endswith("_kv_fp8"):
         return "kv_fp8"
     if name.endswith("_kv_scale"):
@@ -3557,7 +3562,8 @@ def _inferred_edge_lifetime(
         producer_node.op_name == "attention_qkv_projection"
         and consumer_node.op_name
         in {"sparse_mla_fp8_apply", "sparse_mla_fp8_apply_bwd"}
-        and _canonical_path_c_edge_buffer_name(buffer_name) in {"kv_fp8", "kv_scale"}
+        and _canonical_path_c_edge_buffer_name(buffer_name)
+        in {"q_fp8", "q_scale", "kv_fp8", "kv_scale"}
     ):
         return "workspace"
     return "internal"

cppmega_mlx/runtime/path_c_fusion_launcher.py

@@ -53,6 +53,9 @@
     "tl.fusion.physical_abi.physical_buffer_shapes",
 )
 _ROW_DISPATCH_GRID_CHUNKS = "grid_chunks"
+_FORWARD_ONLY_GATE = 0
+_BACKWARD_GATE = 1
+_BACKWARD_FORWARD_PREPASS_GATE = 2
 
 
 def _decode_attr(prim_func: Any, key: str) -> Any:
@@ -98,6 +101,7 @@ class PathCAbiManifest:
     bank_shapes: Mapping[str, tuple[int, ...]]
     bank_dtypes: Mapping[str, str]
     internal_scratch_buffers: tuple[str, ...]
+    internal_scratch_aliases: Mapping[str, Mapping[str, Any]]
     cotangent_seed_buffers: tuple[str, ...]
     top_level_scratch_shapes: Mapping[str, tuple[int, ...]]
     top_level_scratch_dtypes: Mapping[str, str]
@@ -152,6 +156,10 @@ def load_path_c_abi_manifest(prim_func: Any) -> PathCAbiManifest:
     scratch_raw = (
         json.loads(str(scratch_attr)) if scratch_attr is not None else []
     )
+    scratch_alias_attr = prim_func.attrs.get("tl.fusion.internal_scratch_abi_aliases")
+    scratch_aliases = (
+        json.loads(str(scratch_alias_attr)) if scratch_alias_attr is not None else {}
+    )
 
     cotangent_raw = prim_func.attrs.get("tl.fusion.train_step_loss_cotangent_abi")
     cotangent_buffers: tuple[str, ...]
@@ -206,6 +214,7 @@ def load_path_c_abi_manifest(prim_func: Any) -> PathCAbiManifest:
         bank_shapes=bank_shapes,
         bank_dtypes=bank_dtypes,
         internal_scratch_buffers=tuple(scratch_raw),
+        internal_scratch_aliases=scratch_aliases,
         cotangent_seed_buffers=cotangent_buffers,
         top_level_scratch_shapes=top_level_shapes,
         top_level_scratch_dtypes=top_level_dtypes,
@@ -274,8 +283,8 @@ def _flat_into_bank(
             f"expected {placement.dtype}, got {value.dtype}"
         )
     flat = value.reshape((placement.size,))
-    indices = mx.arange(placement.offset, placement.offset + placement.size, dtype=mx.int32)
-    return bank.at[indices].add(flat - bank[indices])
+    slot = slice(placement.offset, placement.offset + placement.size)
+    return bank.at[slot].add(flat - bank[slot])
 
 
 def _bank_slot_view(
@@ -489,6 +498,8 @@ def __call__(
         # (uses top_level_scratch_shapes).
         scratch_buffers: dict[str, mx.array] = {}
         for name in manifest.internal_scratch_buffers:
+            if name in manifest.internal_scratch_aliases:
+                continue
             if name in manifest.logical_to_physical:
                 placement = manifest.logical_to_physical[name]
                 scratch_buffers[name] = mx.zeros(
@@ -545,7 +556,7 @@ def __call__(
                 positional.append(0)
             elif logical_name == manifest.backward_gate_param:
                 backward_gate_param_index = len(positional)
-                positional.append(1 if run_backward else 0)
+                positional.append(_BACKWARD_GATE if run_backward else _FORWARD_ONLY_GATE)
             else:
                 raise ValueError(
                     f"compiled PrimFunc has an unexpected parameter {param_name!r} "
@@ -676,19 +687,16 @@ def launch_chunk(chunk_index: int, subchunk_index: int) -> None:
                 if row_subchunk_param_index is not None:
                     positional[row_subchunk_param_index] = int(subchunk_index)
                 apply_returned_outputs(self._kernel(*positional))
-                eval_positionals()
 
             if run_backward and backward_gate_param_index is not None:
-                positional[backward_gate_param_index] = 0
-                forward_subchunk_count = int(manifest.row_subchunk_count or 1)
+                positional[backward_gate_param_index] = _BACKWARD_FORWARD_PREPASS_GATE
                 for chunk_index in range(chunk_count):
-                    for subchunk_index in range(forward_subchunk_count):
-                        launch_chunk(chunk_index, subchunk_index)
+                    launch_chunk(chunk_index, 0)
                 final_forward_states = snapshot_logical_buffers(
                     _STATE_AFTER_LOGICAL_BUFFERS
                 )
                 restore_logical_buffers(initial_backward_states)
-                positional[backward_gate_param_index] = 1
+                positional[backward_gate_param_index] = _BACKWARD_GATE
                 launch_chunk(0, 0)
                 restore_logical_buffers(final_forward_states)
             else: