Skip to content

RustamSheoran/DWDP-Triton-MoE-Inference-Engine

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

12 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

There is no GPU available offline in this environment, so avoid CUDA-dependent runs here. Python is available for CPU-only checks and documentation work.

Router Module

The repository now includes a production-oriented MoE router package under DWDP/router.

The router is responsible only for expert selection:

  • router logits
  • routing probabilities
  • top-k expert indices
  • normalized routing weights
  • routing metadata

It does not perform dispatch, expert execution, scheduling, communication, or output merging.

Detailed engineering documentation is available in docs/router.md. A package-local overview is available in DWDP/router/README.md.

Dispatcher Module

The repository also includes a production-oriented dispatcher package under DWDP/dispatcher.

The dispatcher consumes completed router output and converts token-major routing assignments into an expert-major physical layout. It is responsible for:

  • expert-major grouping
  • per-expert counts
  • expert offsets
  • token permutation
  • inverse permutation
  • packed token indices
  • packed routing weights
  • reusable dispatch metadata

It does not perform routing, expert execution, communication, scheduling, or output merging.

Detailed engineering documentation is available in docs/dispatcher.md. A package-local overview is available in DWDP/dispatcher/README.md.

Scheduler Module

The scheduler package under DWDP/scheduler consumes DispatchPlan and produces ExecutionPlan.

The scheduler is responsible only for execution planning:

  • expert execution order
  • expert work queues
  • expert-major execution ranges
  • execution priorities
  • stream assignment placeholders
  • dependency metadata placeholders
  • synchronization metadata placeholders
  • scheduler statistics

It does not execute experts, move tensors, launch communication, inspect router output, inspect model weights, or merge outputs.

Detailed engineering documentation is available in docs/scheduler.md. A package-local overview is available in DWDP/scheduler/README.md.

Comms Planner Module

The communication planner package under DWDP/comms_planner consumes ExecutionPlan and produces CommunicationPlan.

The Comms Planner is responsible only for communication planning metadata:

  • local and remote expert classification
  • communication graph metadata
  • transfer descriptors
  • communication groups
  • topology metadata
  • dependency metadata
  • synchronization placeholders
  • prefetch placeholders
  • overlap placeholders
  • communication cost estimates
  • communication statistics

It does not execute communication, move tensors, allocate communication buffers, prefetch weights, execute experts, launch CUDA kernels, launch collectives, or mutate Scheduler output.

Detailed engineering documentation is available in docs/comms_planner.md. A package-local overview is available in DWDP/comms_planner/README.md.

Executor Module

The executor package under DWDP/executor consumes hidden states, DispatchPlan, ExecutionPlan, and CommunicationPlan, then produces ExecutorOutput.

The Executor is responsible only for expert computation:

  • gather token activations for scheduled expert ranges
  • execute expert modules
  • apply routing weights
  • write packed expert outputs
  • emit metadata required by the future Merger

It does not route, dispatch, schedule, plan communication, execute communication, or merge outputs.

Detailed engineering documentation is available in docs/executor.md. A package-local overview is available in DWDP/executor/README.md.

Merger Module

The merger package under DWDP/merger consumes ExecutorOutput and reconstructs the final hidden states for the next Transformer layer.

The Merger is responsible only for output reconstruction:

  • restore token-major assignment order
  • accumulate Top-K expert outputs
  • optionally apply routing weights
  • reshape back to the original token layout
  • emit merge statistics and metadata

It does not route, dispatch, schedule, plan communication, execute communication, execute experts, or inspect upstream runtime plans.

Detailed engineering documentation is available in docs/merger.md. A package-local overview is available in DWDP/merger/README.md.

Runtime Integration Layer

The runtime integration package under DWDP/runtime orchestrates the complete MoE pipeline:

Router -> Dispatcher -> Scheduler -> Comms Planner -> Executor -> Merger

The runtime owns stage modules and reusable workspaces, exposes a HF-style DWDPRuntime API, and provides adapter, profiling, correctness, CLI, and benchmark scaffolding.

Detailed engineering documentation is available in docs/runtime.md. A package-local overview is available in DWDP/runtime/README.md.

Hugging Face Adapter Layer

The adapter package under DWDP/adapters automatically detects supported Hugging Face MoE models and replaces only their MoE blocks with DWDP-backed execution.

The current automatic target is Qwen1.5/Qwen2-style MoE blocks. The adapter preserves native Hugging Face tokenization, attention, KV cache, generation, sampling, checkpoint loading, and non-MoE layers.

Detailed engineering documentation is available in docs/adapters.md. A package-local overview is available in DWDP/adapters/README.md.

Benchmark Reporting

The benchmark reporting package under DWDP/benchmarking provides reproducible experiment directories, structured JSON artifacts, Markdown reports, environment capture, correctness metrics, runtime statistics, and future extension points for long-term performance tracking.

Detailed documentation is available in docs/benchmark_reporting.md.

About

Research-oriented Mixture-of-Experts inference runtime exploring communication-aware scheduling, modular execution pipelines, and distributed GPU systems.

Resources

Stars

1 star

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors