Skip to content

Add exportGraphToDot diagnostic utility#388

Open
arkavo-com wants to merge 1 commit intoml-explore:mainfrom
arkavo-ai:feature/export-graph-to-dot
Open

Add exportGraphToDot diagnostic utility#388
arkavo-com wants to merge 1 commit intoml-explore:mainfrom
arkavo-ai:feature/export-graph-to-dot

Conversation

@arkavo-com
Copy link
Copy Markdown

Summary

Exposes the existing C++ export_to_dot graph visualization function through a Swift API.

let logits = model(input, cache: cache)
exportGraphToDot(path: "/tmp/graph.dot", output: logits)

Why this matters

This utility was essential for discovering a critical performance issue where Swift MLX models silently inject AsType (type cast) nodes from float32 literal defaults. By diffing the DOT graph between Python and Swift, we found:

  • Python graph: 1852 nodes, 0 AsType
  • Swift graph: 2976 nodes, 1046 AsType (from MLXArray.ones() defaulting to float32)

The 1046 extra AsType nodes caused 237 MB of Metal cache churn per token (vs Python's 2.3 MB) and a 12% generation speed regression.

Without this diagnostic, the root cause would have been invisible — the model produced correct output and all individual operation tests showed perfect buffer donation.

Test plan

🤖 Generated with Claude Code

Exposes the C++ export_to_dot function via a Swift-friendly API.
Essential for diagnosing dtype contamination and graph topology
differences between Swift and Python.

Usage:
  let logits = model(input, cache: cache)
  exportGraphToDot(path: "/tmp/graph.dot", output: logits)
  // Then: grep -c 'label ="AsType"' /tmp/graph.dot
  // Non-zero AsType count indicates dtype mismatches

This was critical for discovering that Swift MLX models can
silently inject 1000+ AsType nodes from float32 literal defaults,
causing 100x Metal cache churn vs Python.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
const mlx_imported_function xfunc,
const mlx_vector_array args,
const mlx_map_string_to_array kwargs);
int mlx_export_to_dot_file(const char* path, mlx_array output);
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need this change on the mlx-c side. I suspect you made changes locally to that. it is a good idea, I will talk to Ronan.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants