mirror of
https://github.com/tensorflow/tensorflow.git
synced 2026-10-29 00:15:18 +08:00
Misha Gutman
aaac8ba37b
`false`) and an MoE node's `expert_indices` tensor is a partition input with more than one token (`B * N > 1`), bucket the unique `(token, selected expert)` pairs by expert on the host with a counting sort before invoking the YNNPACK runtime: - `slot_token` `[E, C]` `int32` maps each expert slot to the token it computes, where `C` is the exact maximum number of tokens routed to one expert (padding slots compute token `0`). - `token_slot` `[B, N, K]` `int32` maps each `(token, selected expert)` pair to its slot in `[0, E * C)`. The subgraph dynamically quantizes tokens `[B * N, D_in]` once, gathers them to `[E, C, D_in]`, evaluates the expert FFN with `ynn_define_dot` against 3D `[E, D_in, *]` weights, and gathers each token's `K` expert outputs back by `token_slot` into `[B, N, K, D_out]` before applying routing weights and reducing over `K`. When `dispatch_moe_on_host` is `false`, for single-token decode (`B * N == 1`), or when `expert_indices` is internal to the partition, keep the existing weight-gather lowering. Also allow weight tensors with a per-(output channel, expert) multi-axis quantization that restates their explicit per-expert scale tensor. PiperOrigin-RevId: 1008325688