[NCC_ICFG001] shard_on_i_hybrid MoE kernel fails to compile at seq_len > 2048 (multiple exit nodes in sg0000)

[UbaidHunts]

Summary:
The shard_on_i_hybrid MoE CTE kernel (blockwise_mm_baseline_shard_intermediate_hybrid in nkilib/core/moe/moe_cte/bwmm_shard_on_I.py) compiles successfully at --max-model-len ≤ 2048 but fails at > 2048 with an internal compiler error:

[INTERNAL_ERROR] [NCC_ICFG001] Control flow graph of function sg0000 has multiple exit nodes: Block1_Resume_1_8, Block1_Resume_8

The first (and every) MoE-kernel compile fails — 32/32 compilations error, 0 succeed — so the server never reaches serving. Downstream gloo "Connection closed by peer" / tp_barrier() errors on the other 15 workers are secondary (TP0 dies on the compile failure and the rest time out at the TP barrier).

Enabling shard_on_i_hybrid kernel:
We enabled shard_on_i_hybrid for the gpt-oss-20b BF16 model as a plain, minimal fork (pretty straightforward). We did not modify the kernel. Our model code only calls the stock moe_cte dispatcher (nki_moe_cte_kernel) with implementation=shard_on_i_hybrid, passing standard inputs:

1. conditions of shape [N+1] with a trailing 0 (exactly per the kernel's documented contract),
2. num_static_block=0 (pure-dynamic schedule)

The blockwise_mm_baseline_shard_intermediate_hybrid kernel itself is unmodified stock nkilib. So this failure reproduces through the standard public kernel path with default parameters, no custom kernel authoring on our side.

Reproduction-Launch Command:

vllm serve openai/gpt-oss-20b \
  --tensor-parallel-size 16 \
  --enable-expert-parallel \
  --max-model-len 4096 \
  --max-num-seqs 1 \
  --no-enable-prefix-caching \
  --port 8000 \
  --hf-overrides '{"quantization_config": {}}' \
  --additional-config '{
      "neuron_config": {
          "ep_degree": 4,
          "moe_kernel_mode": "shard_on_i_hybrid",
      }
  }'

We added moe_kernel_mode flag trivially.

Error:
Complete error log file is attached:

error.log

Question

1. Is this a bug with framework or compiler? How do we solve it?
2. Is nl.dynamic_range (hardware dynamic while-loop) known to produce a multi-exit CFG that NCC_ICFG001 rejects at larger sequence lengths?

Environment:

1. NeuronX Compiler version 2.0.262197.0a0+8da1ecb4
2. vLLM evrsion: 0.21.0
3. torch 2.11.0
4. torch-xla version 2.11.0
5. Instance: trn2.48x large
6. Model: openai/gpt-oss-20b and 120b

[github-actions]

Hi @UbaidHunts, Thank you for filing the issue! We will take a look and get back to you.