Gpt_oss prefill: avoid packed expert-axis slicing in chunked MoE to restore MXFP6 constant compression

[abhishek-singh591]

Root cause and fix summary (prefill-only)

Before (problematic pattern)

In QEffPrefillOnlyChunkedGptOssMLP.forward, the prefill path did:

1. Repacked expert tensors into grouped views:
   • W_g/W_u/W_d and biases were reshaped to [num_nsp, local_experts, ...].
2. Sliced that grouped axis inside the local_slot loop:
   • W_g[:, local_slot], W_u[:, local_slot], etc.

This matches the compiler pattern we observed:

• input like float16<16 x 2 x 2880 x 2880>
• then two slices on axis-1 (index 0 and 1)

So even with logically split gate/up params, this path reintroduced a packed-then-sliced representation during subfunction export.

After (fixed pattern)

The updated code now:

1. Keeps routing weights in flat per-expert layout:
   • routing_weights_by_expert = routing_weights.transpose(0, 1) → [num_experts, T]
2. Computes explicit contiguous expert ranges per local_slot:
   • slot_start = local_slot * num_nsp
   • slot_end = slot_start + num_nsp
3. Passes per-slot expert blocks directly:
   • self.experts.gate_proj[slot_start:slot_end]
   • self.experts.up_proj[slot_start:slot_end]
   • self.experts.down_proj[slot_start:slot_end]
   • same pattern for biases and routing weights

So we no longer materialize [num_nsp, local_experts, ...] temporaries and no longer slice that axis inside the subfunction body.

Fix implemented

GPT-OSS chunked MLP now auto-rebuilds blocked tensors whenever expert_blocking_num_nsp changes:

• Added _ensure_blocked_expert_weights(num_nsp) in QEfficient/transformers/models/gpt_oss/modeling_gpt_oss.py.
• Called it at start of forward() before blocked MoE compute.

This removes the need for manual __qeff_init__() calls in tests when only expert-blocking attrs are changed.

[vbaddi]

@abhishek-singh591 what tests are verified on this? Disaggregated mode example?

[abhishek-singh591]

@abhishek-singh591 what tests are verified on this? Disaggregated mode example?

Yes i tested this till compile of prefill+chunking since the changes were only made on the prefill part, didn't check the Perf.

[abhishek-singh591]

QPC size went down from 37Gb to 17GB for full model with 4TS.

Perf with 2 layer:-

Perf before:-
time for this run=0.39685487747192383
time for first run of decode with KV as input = 0.01200103759765625 sec

decode tok/sec=83.26760350742597

Perf after:-
time for this run=0.3967604637145996
time for first run of decode with KV as input = 0.012997150421142578 sec

decode tok/sec=80.18426492513731.

[ochougul]

why is this needed, shouldn't qeff_init always run and do this? is this check really needed?

[abhishek-singh591]

Yes just added in the readme, expert_blocking_num_nsp is set after transform in the test, and GPT-OSS blocked tensors (W_g/W_u/W_d, b_g/b_u/b_d) are shape-dependent on that value.

So, Without re-init, tensors stay in old layout (usually initialized with default num_nsp=num_experts), so matmul/add hits axis mismatch later. The disagg failure (a=16 vs b=32) is exactly that: runtime used num_nsp=16 while blocked params were still built for num_nsp=32.

This removes the need for manual __qeff_init__() calls in tests when only expert-blocking attrs are changed.

[ochougul]

Then, this layout rewrite should be called outside before export, and not in forward method.

[abhishek-singh591]

But that would require updating all the test files. We also need to handle cases where export is called multiple times with expert_blocking_num_nsp, since subsequent export calls could fail due to the modified code. So, this logic should probably live in the model file itself, though not necessarily inside the forward() method.