Skip to content

fix(moe): correct MoE grouped-matmul graphs to match hipDNN contract - #40

Closed
SamuelReeder wants to merge 1 commit into
mainfrom
users/sareeder/fix-moe-grouped-matmul
Closed

fix(moe): correct MoE grouped-matmul graphs to match hipDNN contract#40
SamuelReeder wants to merge 1 commit into
mainfrom
users/sareeder/fix-moe-grouped-matmul

Conversation

@SamuelReeder

Copy link
Copy Markdown
Collaborator

Summary

All MoE grouped-matmul graphs in Workloads/headline/moe.tar.gz and
Workloads/microbench/cudnn_bench_moe.tar.gz failed opgraph-level
deserialization against hipDNN. This corrects the graphs to match hipDNN's
MoE grouped-matmul node contract.

Risk assessment

Low. This change updates two DVC-tracked workload archives and their
internal manifests. It does not change hipDNN or provider source code.

Testing summary

  • tools/check_deserialize.py --level opgraph (built hipDNN, from_json +
    validate + build_operation_graph): 191/191 graphs pass, fail=0.
  • Verified after packaging round-trip (extract from the committed tarballs).
  • Confirmed no MoE forward/backward nodes remain with rank-2 tensors,
    non-INT32 routing tensors, or a non-zero top_k outside SCATTER mode.

Testing checklist

  • Validate opgraph-level deserialization on both tarballs.
  • Validate after packaging round-trip.
  • Run the applicability sweep on gfx942 and gfx950.

Technical changes

  • Reshaped tensors to hipDNN's required rank-3 convention
    ([1, T, K] / [E, K, N] / [E, 1, 1]).
  • Corrected first_token_offset from an E+1-length cumulative-boundary
    array to the E-length per-expert start-offset array hipDNN expects.
  • Zeroed top_k on NONE/GATHER mode nodes (only SCATTER uses it).
  • Removed a token_ks input incorrectly carried by GATHER-mode dgrad
    nodes (forbidden in that mode).
  • Corrected token_index/token_ks tensors to INT32.
  • Split graphs with a disconnected router-GEMM component into separate
    *__router.json files (hipDNN requires one connected component per graph).
  • Corrected SCATTER-mode output row count to the dispatched-token count.
  • Excluded 21 wgrad graphs whose token/doutput row counts mismatch
    (pre- vs post-scatter counts); the current backward node has no
    scatter/token_index support, so this isn't a shape fix. Documented as a
    known gap in the manifest.

All MoE graphs in moe.tar.gz and cudnn_bench_moe.tar.gz failed opgraph
deserialization. Root causes:

- Tensors were flattened; hipDNN's MoE node requires rank-3 shapes
  ([1, T, K] / [E, K, N] / [E, 1, 1]).
- first_token_offset used an E+1-length cumulative-boundary array;
  hipDNN expects the E-length per-expert start-offset array.
- top_k was set to the real routing value on NONE/GATHER nodes, but
  those modes require top_k == 0 (SCATTER only).
- GATHER-mode dgrad nodes carried a token_ks input inherited from
  their paired SCATTER forward node; GATHER forbids token_ks.
- token_index/token_ks used the activation dtype instead of INT32.
- Some graphs had a disconnected router-GEMM component; hipDNN
  requires one connected component per graph, so these were split
  into separate *__router.json files.
- SCATTER-mode output row count used the reduced final-token count;
  hipDNN's contract requires the dispatched-token count.

Excluded 21 wgrad graphs whose token/doutput row counts mismatch
(pre- vs post-scatter counts) -- MoeGroupedMatmulBwdAttributes has no
scatter/token_index support today, so this isn't fixable by reshaping.

Result: 191/191 graphs (165 moe.tar.gz + 26 cudnn_bench_moe.tar.gz)
pass from_json + validate + build_operation_graph.
@SamuelReeder

Copy link
Copy Markdown
Collaborator Author

Superseded — the moe fix is now part of #39 (norm + moe on one branch), no need for a separate PR for these DVC-only changes.

@SamuelReeder
SamuelReeder deleted the users/sareeder/fix-moe-grouped-matmul branch September 1, 2026 19:52
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.

1 participant