Skip to content

[Levanter] Add DeepEP profiling scopes#6213

Open
dlwh wants to merge 1 commit into
mainfrom
codex/6139-deepep-trace-scopes
Open

[Levanter] Add DeepEP profiling scopes#6213
dlwh wants to merge 1 commit into
mainfrom
codex/6139-deepep-trace-scopes

Conversation

@dlwh

@dlwh dlwh commented Jun 6, 2026

Copy link
Copy Markdown
Member

Add named scopes around DeepEP local packing, dispatch transport, and combine transport so profiler traces can attribute movement and bookkeeping separately. This is instrumentation only and does not change the MoE data path.

Part of #6139

@dlwh dlwh added the agent-generated Created by automation/agent label Jun 6, 2026
@claude

claude Bot commented Jun 6, 2026

Copy link
Copy Markdown
Contributor

Claude finished @dlwh's task in 2m 35s —— View job


Code review

No issues found. Checked for bugs and CLAUDE.md compliance.

The change is pure instrumentation — wraps existing code in jax.named_scope(...) blocks (deepep_pack_local_assignments, deepep_collapse_local_assignments, deepep_layout, deepep_dispatch_transport, deepep_combine_transport). No logic changes; variables defined inside scopes remain accessible to surrounding code; indentation and nesting are correct.

  • Check PR state and prior reviews
  • Locate relevant CLAUDE.md/AGENTS.md files (AGENTS.md, lib/levanter/AGENTS.md)
  • Get PR summary
  • Run review agents (bug + CLAUDE.md compliance)
  • Validate flagged issues (none flagged)
  • Post review
    • Branch: codex/6139-deepep-trace-scopes

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

Labels

agent-generated Created by automation/agent

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant