Wire QEffGptOssMLP into shared MoEBlockMixin - #1312
Open
Shrubabati7 wants to merge 2 commits into
Open
Shrubabati7 wants to merge 2 commits into
Shrubabati7 wants to merge 2 commits into
Conversation
Comment on lines
4
to
-374
| @@ -88,180 +87,7 @@ def transform_weights(self) -> MoEWeights: | |||
| return self.moe_weights | |||
|
|
|||
|
|
|||
| class _QEffGptOssLegacyBlockedMixin: | |||
| def forward(self, hidden: torch.Tensor): | |||
| if os.environ.get("NUM_FFN_BLOCKS", None) is not None: | |||
| return self.blocked_ffn_forward(hidden) | |||
| return super().forward(hidden) | |||
|
|
|||
| def blocked_ffn_forward(self, hidden: torch.Tensor): | |||
| if not getattr(self, "weights_transformed", False): | |||
| raise RuntimeError(f"{type(self).__name__} weights are not transformed; run OptimizedMoEWeightsTransform") | |||
| B, S, H = hidden.shape | |||
| T = B * S | |||
| hidden = hidden.view(T, H) | |||
| weights = self.moe_weights | |||
|
|
|||
| # Router computation | |||
| router_logits = F.linear(hidden, self.router.weight, self.router.bias) | |||
|
|
|||
| # Top-k selection | |||
| top_w, top_i = torch.topk(router_logits, self.router.top_k, dim=-1) # both [T, K] | |||
| top_w = torch.nn.functional.softmax(top_w, dim=1, dtype=top_w.dtype) | |||
|
|
|||
| masked_logits = torch.zeros_like(router_logits) | |||
| masked_logits.scatter_(1, top_i, top_w) | |||
|
|
|||
| # Routing weights for each expert [T, E] | |||
| routing_weights = masked_logits | |||
|
|
|||
| # ────────────────── allocate the output tensor ───── | |||
| expert_out = hidden.new_zeros((T, H)) # accumulation buffer | |||
| target_blocks = int(os.environ.get("NUM_FFN_BLOCKS", 1)) | |||
| block_positions = [] | |||
| for j in range(target_blocks): | |||
| block_positions.append(j * (T // target_blocks)) | |||
| # ───────────────────────── Expert computation loop ───────────────────────────── | |||
| for e in range(self.experts.num_experts): | |||
| routing_weight = routing_weights[:, e].unsqueeze(-1) # [T, 1] | |||
|
|
|||
| W_g, W_u = weights.gate[e], weights.up[e] # [H, I], [H, I] | |||
| b_g, b_u = weights.gate_bias[e], weights.up_bias[e] # [I], [I] | |||
| W_d = weights.down[e] # [I, H] | |||
| b_d = weights.down_bias[e] # [H] | |||
|
|
|||
| block_count = 0 | |||
| outs = [] | |||
| for block_idx in range(target_blocks): | |||
| block_count += 1 | |||
| qi = block_positions[block_idx] | |||
|
|
|||
| # Calculate block size (last block should be handled with remainder) | |||
| if block_idx == target_blocks - 1: | |||
| real_q_len = T - qi | |||
| else: | |||
| real_q_len = block_positions[block_idx + 1] - qi | |||
|
|
|||
| tgb = hidden[qi : qi + real_q_len, :] | |||
| # Gate and Up projections | |||
| # Gate and Up projections | |||
| gate = (tgb @ W_g) + b_g # [T, I] | |||
| up = (tgb @ W_u) + b_u # [T, I] | |||
|
|
|||
| # Apply GptOss activation with clamping | |||
| gate = gate.clamp(min=torch.finfo(torch.float16).min, max=self.experts.limit) | |||
| up = up.clamp(min=-self.experts.limit, max=self.experts.limit) | |||
|
|
|||
| # GLU activation | |||
| glu = gate * torch.sigmoid(gate * self.experts.alpha) | |||
| intermediate = (up + 1) * glu # [T, I] | |||
|
|
|||
| # Down projection | |||
| down_out_block = (intermediate @ W_d) + b_d # [T, H] | |||
|
|
|||
| outs.append(down_out_block) | |||
|
|
|||
| down_out = torch.cat(outs, dim=0) | |||
|
|
|||
| # Apply routing weights and accumulate | |||
| expert_out += down_out * routing_weight | |||
|
|
|||
| # original shape [B, S, H] | |||
| return expert_out.view(B, S, H), router_logits | |||
|
|
|||
| def blocked_ffn_forward_block_weights(self, hidden: torch.Tensor): | |||
| if not getattr(self, "weights_transformed", False): | |||
| raise RuntimeError(f"{type(self).__name__} weights are not transformed; run OptimizedMoEWeightsTransform") | |||
| B, S, H = hidden.shape | |||
| T = B * S | |||
| hidden = hidden.view(T, H) | |||
| weights = self.moe_weights | |||
|
|
|||
| # Router computation | |||
| router_logits = F.linear(hidden, self.router.weight, self.router.bias) | |||
|
|
|||
| # Top-k selection | |||
| top_w, top_i = torch.topk(router_logits, self.router.top_k, dim=-1) # both [T, K] | |||
| top_w = torch.nn.functional.softmax(top_w, dim=1, dtype=top_w.dtype) | |||
|
|
|||
| masked_logits = torch.zeros_like(router_logits) | |||
| masked_logits.scatter_(1, top_i, top_w) | |||
|
|
|||
| # Routing weights for each expert [T, E] | |||
| routing_weights = masked_logits | |||
|
|
|||
| # ────────────────── allocate the output tensor ───── | |||
| expert_out = hidden.new_zeros((T, H)) # accumulation buffer | |||
| target_blocks = int(os.environ.get("NUM_BLOCKS", 1)) | |||
| block_positions = [] | |||
| for j in range(target_blocks): | |||
| block_positions.append(j * (T // target_blocks)) | |||
| # ───────────────────────── Expert computation loop ───────────────────────────── | |||
| for e in range(self.experts.num_experts): | |||
| routing_weight = routing_weights[:, e].unsqueeze(-1) # [T, 1] | |||
|
|
|||
| W_g, W_u = weights.gate[e], weights.up[e] # [H, I], [H, I] | |||
| b_g, b_u = weights.gate_bias[e], weights.up_bias[e] # [I], [I] | |||
| W_d = weights.down[e] # [I, H] | |||
| b_d = weights.down_bias[e] # [H] | |||
|
|
|||
| block_count = 0 | |||
| outs = [] | |||
| for block_idx in range(target_blocks): | |||
| block_count += 1 | |||
| qi = block_positions[block_idx] | |||
|
|
|||
| # Calculate block size (last block should be handled with remainder) | |||
| if block_idx == target_blocks - 1: | |||
| real_q_len = T - qi | |||
| else: | |||
| real_q_len = block_positions[block_idx + 1] - qi | |||
|
|
|||
| tgb = hidden[qi : qi + real_q_len, :] | |||
| # Gate and Up projections | |||
|
|
|||
| wg_col_shape = W_g.shape[1] | |||
| wg_num_blocks = math.ceil(wg_col_shape / 128) | |||
| last_block_size = wg_col_shape % 128 if wg_col_shape % 128 != 0 else 128 | |||
|
|
|||
| intermediates = [] | |||
| for i in range(wg_num_blocks): | |||
| if i == wg_num_blocks - 1: | |||
| cur_gate = (tgb @ W_g[:, -last_block_size:]) + b_g[-last_block_size:] | |||
| cur_up = (tgb @ W_u[:, -last_block_size:]) + b_u[-last_block_size:] | |||
| else: | |||
| cur_gate = (tgb @ W_g[:, i * 128 : (i + 1) * 128]) + b_g[i * 128 : (i + 1) * 128] | |||
| cur_up = (tgb @ W_u[:, i * 128 : (i + 1) * 128]) + b_u[i * 128 : (i + 1) * 128] | |||
|
|
|||
| cur_gate = cur_gate.clamp(min=torch.finfo(torch.float16).min, max=self.experts.limit) | |||
| cur_up = cur_up.clamp(min=-self.experts.limit, max=self.experts.limit) | |||
| cur_glu = cur_gate * torch.sigmoid(cur_gate * self.experts.alpha) | |||
| cur_intermediate = (cur_up + 1) * cur_glu | |||
| intermediates.append(cur_intermediate) | |||
|
|
|||
| intermediate = torch.cat(intermediates, dim=-1) | |||
|
|
|||
| downs = [] | |||
| for i in range(wg_num_blocks): | |||
| if i == wg_num_blocks - 1: | |||
| downs.append((intermediate @ W_d[:, -last_block_size:]) + b_d[-last_block_size:]) | |||
| else: | |||
| downs.append((intermediate @ W_d[:, i * 128 : (i + 1) * 128]) + b_d[i * 128 : (i + 1) * 128]) | |||
|
|
|||
| down_out_block = torch.cat(downs, dim=1) | |||
| outs.append(down_out_block) | |||
|
|
|||
| down_out = torch.cat(outs, dim=0) | |||
|
|
|||
| # Apply routing weights and accumulate | |||
| masked_down = torch.where(routing_weight > 0, down_out * routing_weight, torch.zeros_like(expert_out)) | |||
| expert_out += masked_down | |||
|
|
|||
| # original shape [B, S, H] | |||
| return expert_out.view(B, S, H), router_logits | |||
|
|
|||
|
|
|||
| class QEffGptOssMLP(_QEffGptOssLegacyBlockedMixin, QEffMoEBlockMixin, GptOssMLP): | |||
| class QEffGptOssMLP(QEffMoEBlockMixin, GptOssMLP): | |||
| _moe_return_router_logits = True | |||
| supported_moe_flavours = ( | |||
| MoEFlavour.SIMPLE_LOOP, | |||
| @@ -286,164 +112,12 @@ def moe_profile(self) -> MoEProfile: | |||
|
|
|||
| def route(self, x: torch.Tensor): | |||
| router_logits = F.linear(x, self.router.weight, self.router.bias) | |||
| top_w, top_i = torch.topk(router_logits, self.router.top_k, dim=-1) | |||
| top_w = F.softmax(top_w, dim=1, dtype=top_w.dtype) | |||
| probs = F.softmax(router_logits, dim=-1, dtype=torch.float) | |||
| top_w, top_i = torch.topk(probs, self.router.top_k, dim=-1) | |||
| top_w = top_w / top_w.sum(dim=-1, keepdim=True) | |||
| top_w = top_w.to(x.dtype) | |||
| return (top_i, top_w), router_logits | |||
|
|
|||
| # ------------------- Gather based, weights as activation approach --------------- | |||
| def forward_weights_as_activation(self, hidden_states): | |||
| if not getattr(self, "weights_transformed", False): | |||
| raise RuntimeError(f"{type(self).__name__} weights are not transformed; run OptimizedMoEWeightsTransform") | |||
| bs, seq_len, _ = hidden_states.shape | |||
| hidden_states = hidden_states.view(bs * seq_len, self.experts.hidden_size) | |||
| weights = self.moe_weights | |||
|
|
|||
| # Router computation | |||
| router_logits = F.linear(hidden_states, self.router.weight, self.router.bias) | |||
| router_top_value, router_indices = torch.topk(router_logits, self.router.top_k, dim=-1) | |||
| router_top_value = torch.nn.functional.softmax(router_top_value, dim=1, dtype=router_top_value.dtype) | |||
|
|
|||
| # GATHER - collect weights for selected experts | |||
| gate_proj = weights.gate[router_indices.flatten()] | |||
| up_proj = weights.up[router_indices.flatten()] | |||
| gate_proj_bias = weights.gate_bias[router_indices.flatten()] | |||
| up_proj_bias = weights.up_bias[router_indices.flatten()] | |||
| down_proj = weights.down[router_indices.flatten()] | |||
| down_proj_bias = weights.down_bias[router_indices.flatten()] | |||
|
|
|||
| # Apply Chosen Experts (without routing weights first) | |||
| # expert_in = hidden_states.repeat_interleave(self.router.top_k, dim=0) | |||
| # expert_in = expert_in.view(-1, 1, self.experts.hidden_size) | |||
| # Reshape for bmm: (bs*seq_len*top_k, 1, hidden_size) | |||
| expert_in = ( | |||
| hidden_states.unsqueeze(1) | |||
| .expand(-1, self.router.top_k, -1) | |||
| .contiguous() | |||
| .view(-1, 1, self.experts.hidden_size) | |||
| ) | |||
|
|
|||
| gate = torch.bmm(expert_in, gate_proj) + gate_proj_bias.unsqueeze(1) | |||
| up = torch.bmm(expert_in, up_proj) + up_proj_bias.unsqueeze(1) | |||
|
|
|||
| # Apply activation with clamping | |||
| gate = gate.clamp(min=None, max=self.experts.limit) | |||
| up = up.clamp(min=-self.experts.limit, max=self.experts.limit) | |||
| glu = gate * torch.sigmoid(gate * self.experts.alpha) | |||
| gated_output = (up + 1) * glu | |||
|
|
|||
| experts_out = torch.bmm(gated_output, down_proj) + down_proj_bias.unsqueeze(1) | |||
| experts_out = experts_out.view(bs * seq_len, self.router.top_k, self.experts.hidden_size) | |||
|
|
|||
| # Apply routing weights AFTER expert computation (This is before on Llama4) | |||
| experts_out = experts_out * router_top_value.unsqueeze(-1) | |||
| experts_out = torch.einsum("bnd->bd", experts_out) | |||
|
|
|||
| return experts_out, router_logits | |||
|
|
|||
| def forward(self, hidden_states): | |||
| if os.environ.get("NUM_FFN_BLOCKS", None) is not None: | |||
| return self.blocked_ffn_forward(hidden_states) | |||
| return QEffMoEBlockMixin.forward(self, hidden_states) | |||
|
|
|||
| def optimized_moe_forward(self, hidden_states: torch.Tensor): | |||
| if not getattr(self, "weights_transformed", False): | |||
| raise RuntimeError(f"{type(self).__name__} weights are not transformed; run OptimizedMoEWeightsTransform") | |||
| B, S, H = hidden_states.shape | |||
| T = B * S | |||
| hidden_states = hidden_states.view(T, H) | |||
| weights = self.moe_weights | |||
|
|
|||
| # Router computation | |||
| router_logits = F.linear(hidden_states, self.router.weight, self.router.bias) | |||
|
|
|||
| # Top-k selection | |||
| top_w, selected_experts = torch.topk(router_logits, self.router.top_k, dim=-1) # both [T, K] | |||
| top_w = torch.nn.functional.softmax(top_w, dim=1, dtype=top_w.dtype) | |||
|
|
|||
| # Creating experts mask and routing weights masked | |||
| awesome_experts_mask_1 = ( | |||
| torch.nn.functional.one_hot(selected_experts[:, 0], num_classes=self.experts.num_experts) | |||
| .bool() | |||
| .T.unsqueeze(-1) | |||
| ) | |||
| awesome_experts_mask_2 = ( | |||
| torch.nn.functional.one_hot(selected_experts[:, 1], num_classes=self.experts.num_experts) | |||
| .bool() | |||
| .T.unsqueeze(-1) | |||
| ) | |||
| awesome_experts_mask_3 = ( | |||
Contributor
There was a problem hiding this comment.
This removes the only GPT-OSS runtime use of NUM_FFN_BLOCKS: the shared MoE path does not consume it, and the new MoE config options are not a compatibility mapping for its token-blocking behavior.
If removing this legacy control is intentional, please also remove its stale handling from modeling_auto.py (export sizing/hash) and update examples/disagg_serving/README.md and if backward compatibility is needed put the warning as well.
Author
There was a problem hiding this comment.
Restored: QEffGptOssMLP.execute_moe_flavour now overrides the flavour dispatch unconditionally when NUM_FFN_BLOCKS is set (matching the old mixin behaviour), modeling_auto.py hash/sizing logic retained with docstring, and README.md updated to reflect the token-blocking semantics.
| module.expert_blocking_packed_chunk_size = expert_parallel_chunk_size | ||
| expert_intermediate_block_size = moe_config.get("expert_intermediate_block_size", None) | ||
| if expert_intermediate_block_size is not None: | ||
| expert_intermediate_block_size = int(expert_intermediate_block_size) |
Contributor
There was a problem hiding this comment.
We can put a check for negative values and reject non-positive values, as this value is passed through int() before being used as the tile-loop step. This is mainly for robustness.
Author
There was a problem hiding this comment.
Change done: added if expert_intermediate_block_size <= 0: raise ValueError(...) immediately after the int() cast in OptimizedMoEExportConfigTransform.
| packed_chunk_size = max(1, min(seq_len // num_packed_chunks, seq_len)) | ||
| matched_idx = build_matched_idx_from_cumsum(T2Ei) | ||
| valid_rows = torch.einsum("ij->i", T2Ei.to(torch.int32)).unsqueeze(1) | ||
| valid_rows = T2Ei.to(torch.int32).sum(dim=1, keepdim=True) |
Contributor
There was a problem hiding this comment.
.sum() will fail on current mainline, did you validate this?
| f"seq_len={seq_len} must be divisible by num_packed_chunks={num_packed_chunks}" | ||
| ) | ||
| packed_chunk_size = seq_len // num_packed_chunks | ||
| packed_chunk_size = max(1, min(seq_len // num_packed_chunks, seq_len)) |
Contributor
There was a problem hiding this comment.
We don't need this pls remove, since num_packed_chunks is computed during the MoE transform and the export sequence length is subsequently rounded to a multiple of it in modeling_auto.py.
Shrubabati7
force-pushed
the
gpt_oss_moe_transforms
branch
3 times, most recently
from
September 15, 2026 03:30
516d9fa to
8084de8
Compare
abhishek-singh591
approved these changes
Sep 15, 2026
added 2 commits
September 15, 2026 11:18
…, EXPERT_PARALLEL) using gptoss_clamped_glu_mlp profile. Drop bespoke prefill/decode MoE implementations; add expert_intermediate_block_size support and TTFT reporting. Switch einsum reductions to .sum() following 2dae7e3. Signed-off-by: Shrubabati Moitra <moitra@qti.qualcomm.com>
- route(): topk raw logits first then softmax over top-k only, matching HF GptOssTopKRouter (was softmax-all then topk then redundant renorm) - flavours.py: remove unused chunk_rows variable (ruff F841) - ruff format pass Signed-off-by: Shrubabati Moitra <moitra@qti.qualcomm.com>
vbaddi
force-pushed
the
gpt_oss_moe_transforms
branch
from
September 15, 2026 05:48
d618830 to
439b90d
Compare
Contributor
|
CI-Ready |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.