Feat/moe nsp blocking all models

[divytrip3005]

NSP-Blocked MoE Prefill Dispatch

What Is This?

Large MoE (Mixture-of-Experts) models like Qwen3-MoE, GPT-OSS, GraniteMoE, and Qwen3-VL-MoE route each input token to a small subset of expert FFN networks (e.g. 8 out of 128 experts). During prefill, all tokens are processed simultaneously — this is compute-bound and benefits from parallelism.

On Qualcomm AI100 hardware, the chip has 16 NSPs that can execute in parallel. The standard per-expert sequential loop underutilizes this parallelism. NSP blocking restructures the expert dispatch so that each NSP handles a dedicated group of experts simultaneously, dramatically improving prefill throughput.

---

How It Works

Standard MoE Dispatch (Sequential)

Tokens [T, H]
    │
    ▼
Router → selects top-k experts per token
    │
    ▼
for expert_e in range(num_experts):          # sequential loop over all E experts
    mask = tokens routed to expert_e
    out += expert_e(tokens[mask]) * weight
    │
    ▼
Output [T, H]

Every expert is processed one at a time.

NSP-Blocked Dispatch (Parallel)

Tokens [T, H]
    │
    ▼
Router → routing_weights [T, E]
    │
    ▼
Reshape experts into NSP groups:
  E experts → num_nsp groups × local_experts each
  e.g. 128 experts, NSP=16 → 16 groups × 8 experts

    NSP-0    NSP-1    NSP-2  ...  NSP-15
   [E0..E7] [E8..E15] [E16..E23]  [E120..E127]
      │         │         │              │
      ▼         ▼         ▼              ▼
   (parallel execution across all 16 NSPs)

for slot in range(local_experts):            # loop over 8 slots (not 128!)
    T2Ei = tokens routed to this slot's expert in each NSP group
    │
    ├─ cumsum → compute packed indices of active tokens
    ├─ gather → pack only active tokens [T_active, H]
    ├─ matmul → run expert FFN on packed tokens
    └─ scatter → write results back to output buffer
    │
    ▼
Output [T, H]  (sum across NSP dimension)

The inner loop runs local_experts = E / num_nsp times (e.g. 8 instead of 128), with all 16 NSPs working in parallel on their expert group.

---

Implementation

Key Components

1. Weight Splitting (__qeff_init__)

The fused gate_up_proj [E, H, 2I] is split into separate gate and up projections and reshaped for NSP-parallel access:

def __qeff_init__(self):
    W_gate_up = self.experts.gate_up_proj  # [E, H, 2I]
    I = W_gate_up.shape[2] // 2
    self._W_g = nn.Parameter(W_gate_up[:, :, :I].contiguous())  # [E, H, I]
    self._W_u = nn.Parameter(W_gate_up[:, :, I:].contiguous())  # [E, H, I]
    self._W_d = nn.Parameter(self.experts.down_proj.contiguous())  # [E, I, H]

2. NSP-Grouped Reshape

Experts are grouped across NSPs using a double-transpose pattern that correctly interleaves experts:

local_experts = num_experts // num_nsp  # e.g. 128 // 16 = 8

# routing weights: [T, E] → [num_nsp, local_experts, T]
rw = routing_weights.transpose(0,1).contiguous()
     .view(local_experts, num_nsp, T).transpose(0,1).contiguous()

# weights: [E, H, I] → [num_nsp, local_experts, H, I]
W_g = self._W_g.view(local_experts, num_nsp, H, I).transpose(0,1).contiguous()

3. Cumsum-Scatter-Gather Kernel

For each slot, only tokens routed to that expert group are processed:

for slot in range(local_experts):
    routing_weight = rw[:, slot, :]      # [num_nsp, T]
    T2Ei = routing_weight > 0            # boolean mask of active tokens

    # Pack active tokens using cumsum-based indexing
    matched_idx = _build_matched_idx_from_cumsum(T2Ei)

    # Gather: pick only active tokens → [num_nsp, T_active, H]
    x_chunk = CtxGatherFunc3DGeneralized.apply(x_expanded, matched_idx)

    # Expert FFN on packed tokens
    gate_prime = x_chunk @ W_g[:, slot]
    up_prime   = x_chunk @ W_u[:, slot]
    down_chunk = (up_prime * act_fn(gate_prime)) @ W_d[:, slot]

    # Scatter: write results back
    expert_out = CtxScatterFunc3DGeneralized.apply(expert_out, matched_idx, down_chunk * rw_chunk)

4. Forward Dispatch

The class attribute supports_moe_prefill_blocking = True signals that this module supports API-driven config injection. The forward method dispatches to the NSP path only when expert_blocking_num_nsp is set:

class QEffPrefillChunkedQwen3VLMoeTextSparseMoeBlock(Qwen3VLMoeTextSparseMoeBlock):
    supports_moe_prefill_blocking = True

    def forward(self, hidden_states):
        ...
        if hasattr(self, "expert_blocking_num_nsp"):
            expert_out = self._forward_expert_blocked(x, routing_weights)
            return expert_out.view(B, S, H), router_logits
        return self.orig_forward(hidden_states)  # fallback

---

Supported Models

Model	Class	Experts	NSP Groups (NSP=16)
Qwen3-MoE	QEffPrefillChunkedQwen3MoeSparseMoeBlock	128	16 × 8
GPT-OSS	QEffPrefillOnlyChunkedGptOssMLP	128	16 × 8
GraniteMoE	QEffPrefillChunkedGraniteMoeMoE	32	16 × 2
Qwen3-VL-MoE	QEffPrefillChunkedQwen3VLMoeTextSparseMoeBlock	128	16 × 8

---

How to Enable

CausalLM Models (Qwen3-MoE, GPT-OSS, GraniteMoE)

from QEfficient import QEFFAutoModelForCausalLM

model = QEFFAutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B-Instruct-2507")

# Compile prefill with NSP blocking enabled
prefill_qpc = model.compile(
    batch_size=1,
    prefill_seq_len=512,
    ctx_len=4096,
    num_cores=16,                        # number of NSPs on the chip
    num_devices=1,
    mxfp6_matmul=True,
    mxint8_kv_cache=True,
    prefill_only=True,
    enable_chunking=True,
    moe_prefill_num_nsp=16,              # ← enable NSP blocking (set to num_cores)
    moe_prefill_packed_chunk_size=256,   # ← token chunk size per iteration
)

# Compile decode (no blocking needed for decode)
decode_qpc = model.compile(
    batch_size=1,
    prefill_seq_len=1,
    ctx_len=4096,
    num_cores=16,
    num_devices=1,
    mxfp6_matmul=True,
    mxint8_kv_cache=True,
)

VLM Models (Qwen3-VL-MoE)

from QEfficient import QEFFAutoModelForImageTextToText

model = QEFFAutoModelForImageTextToText.from_pretrained(
    "Qwen/Qwen3-VL-30B-A3B-Instruct",
    attn_implementation="eager",
    kv_offload=True,
)

# Compile prefill with NSP blocking
prefill_qpc = model.compile(
    batch_size=1,
    prefill_seq_len=512,
    ctx_len=4096,
    height=354,
    width=536,
    num_cores=16,
    num_devices=1,
    mxfp6_matmul=True,
    mxint8_kv_cache=True,
    prefill_only=True,
    enable_chunking=True,
    skip_vision=True,                    # compile lang model only
    moe_prefill_num_nsp=16,              # ← enable NSP blocking
    moe_prefill_packed_chunk_size=256,
)

# Compile decode
decode_qpc = model.compile(
    batch_size=1,
    prefill_seq_len=1,
    ctx_len=4096,
    height=354,
    width=536,
    num_cores=16,
    num_devices=1,
    mxfp6_matmul=True,
    mxint8_kv_cache=True,
    prefill_only=False,
    skip_vision=True,
)

Disabling NSP Blocking (Baseline)

To compile without NSP blocking for comparison:

prefill_qpc = model.compile(
    ...,
    prefill_only=True,
    enable_chunking=True,
    moe_prefill_num_nsp=None,   # ← None disables NSP blocking
)

---

Key Parameters

Parameter	Type	Default	Description
moe_prefill_num_nsp	Optional[int]	None	Number of NSP groups to split experts across. Set to num_cores (typically 16) to enable blocking. None disables blocking.
moe_prefill_packed_chunk_size	int	256	Number of token rows per packed chunk in the scatter-gather loop. Controls the tradeoff between loop iterations and chunk size.

---

Notes

• Token correctness: NSP blocking produces numerically identical outputs to the standard path for shallow models (2 layers). For full-depth models, minor FP16 rounding differences may accumulate across layers due to non-associativity of floating point addition — this is expected and does not affect model quality in practice.
• Cache isolation: Different values of moe_prefill_num_nsp produce different ONNX/QPC hashes, so switching between blocking configs does not cause cache conflicts.
• Decode path: NSP blocking applies only to the prefill phase. The decode path uses a standard index_select-based dispatch which is already efficient for single-token generation.

[divytrip3005]

@vbaddi can you review this PR ?

[vbaddi]

There is ongoing effort to put all the changes in 1.22_tmp branch which should soon come up in mainline QEff. Let's rebase once that is done to streamline this work item. thanks @divytrip3005

cc: @quic-rishinr @anujgupt-github

[vbaddi]

nit: no need of this imo, this would be same as num_cores already.

[vbaddi]

nit: Let's rebase to 1.22_tmp, these changes should catch up in there.

[vbaddi]

nit: what's the purpose of this? why do we need different chunked attention? does this solve anything unique?

[vbaddi]

nit: let's rebase to 1.22_tmp branch, this logic should be aligned w/gptoss and qwen3-moe.

[vbaddi]

nit: what changed here? why so much diff? what this formatting issue earlier?

[vbaddi]

nit: how are we verifying the chunked export here? should be either via customops (CtxGeneralized*) in onnx or checking the module presence in pytorch no?