Skip to content

feature: add routed MoE emulator substrate#91

Open
Happymic wants to merge 4 commits into
mainfrom
feat/routed-moe-emulator-substrate
Open

feature: add routed MoE emulator substrate#91
Happymic wants to merge 4 commits into
mainfrom
feat/routed-moe-emulator-substrate

Conversation

@Happymic

Copy link
Copy Markdown
Collaborator

Summary

Adds the transactional-emulator substrate for routed MoE execution.

Changes

  • Add dynamic routing and expert-weight addressing support.
  • Add vector top-k softmax execution and BF16 router-path support.
  • Add vector scalar min/max execution for GPT-OSS clamp-gated expert activation.
  • Add gather/scatter and HBM/VRAM layout helpers for routed expert execution.
  • Add GPT-OSS staged MoE correctness tests for gate/up, clamp, activation, combine, router, top-k, gather/scatter, and real-layer validation.
  • Add attention/glue semantic checks used by the routed MoE validation path.

Scope notes

Validation

Ran locally from transactional_emulator:

cargo fmt --check
cargo test

@Happymic Happymic force-pushed the feat/routed-moe-emulator-substrate branch 2 times, most recently from f285d19 to 79b3b2b Compare June 27, 2026 01:21
@Happymic Happymic force-pushed the feat/routed-moe-emulator-substrate branch from 79b3b2b to 8a40efc Compare June 27, 2026 01:25
Comment on lines +30 to +38
unsafe {
tch::Tensor::from_blob(
data.as_ptr() as *const u8,
&[data.len() as i64],
&[],
tch::Kind::Float,
tch::Device::Cpu,
)
.internal_to_copy((tch::Kind::Float, tch::Device::Cpu), false)

@qichao-arlo-wang qichao-arlo-wang Jun 29, 2026

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.

Why unsafe here?

Comment on lines +282 to +333
fn classify_comment(comment: &str, current: StageKind) -> StageKind {
let text = comment.to_ascii_lowercase();
if text.contains("gpt-oss router")
|| text.contains("router token")
|| text.contains("router dot token")
{
StageKind::RouterTopk
} else if text.contains("gpt-oss vram scatter-add") || text.contains("_scatter") {
StageKind::ScatterCombine
} else if text.contains("allocate vram matrix step6_pair") && text.contains("_route") {
StageKind::ExpertRouteWeight
} else if text.contains("materialize route weight")
|| text.contains("vram matrix mul")
|| (text.contains("true-zero vram rows") && matches!(current, StageKind::ExpertRouteWeight))
{
StageKind::ExpertRouteWeight
} else if text.contains("step6_device_routing_acc") || text.contains("true-zero vram rows") {
StageKind::AccumulatorInit
} else if text.contains("gpt-oss gather token rows")
|| text.contains("gather pair")
|| text.contains("clear gather padding")
|| (text.contains("allocate vram matrix step6_pair") && text.contains("_gather"))
{
StageKind::Gather
} else if text.contains("dynamic expert bias add") {
StageKind::ExpertBias
} else if text.contains("allocate vram matrix step6_pair") && text.contains("_sigmoid") {
StageKind::ExpertActivation
} else if text.contains("tile row min fp")
|| text.contains("tile row max fp")
|| matches!(current, StageKind::ExpertActivation)
&& (text.contains("vram fill zero")
|| text.contains("vram matrix add")
|| text.contains("vram matrix mul"))
{
StageKind::ExpertActivation
} else if text.contains("dynamic hbm weight prefetch")
|| text.contains("expert_id_to_weight_base")
{
StageKind::ExpertWeightAddress
} else if text.contains("subblock [") {
StageKind::ExpertWeightPrefetch
} else if text.contains("sub projection")
|| text.contains("vram block add")
|| text.contains("vram block")
|| (text.contains("allocate vram matrix step6_pair") && !text.contains("_gather"))
{
StageKind::ExpertProjection
} else {
current
}
}

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.

nasty code

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