[RunkaiTao]

Purpose

In #32655, we used heuristic condition to decide whether turning on unpermute or not. In this PR, we also support unpermute control from configurations directly. We can see some improvement when adding naive_block_assignment in the search space during tuning.

Test Result

gpt-oss 120b
1600 input, 600 output, max_loras = 8, concurrency = 8

before

============ Serving Benchmark Result ============
Successful requests:                     200       
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  206.03    
Total input tokens:                      316188    
Total generated tokens:                  119989    
Request throughput (req/s):              0.97      
Output token throughput (tok/s):         582.38    
Peak output token throughput (tok/s):    688.00    
Peak concurrent requests:                16.00     
Total token throughput (tok/s):          2117.04   
---------------Time to First Token----------------
Mean TTFT (ms):                          562.44    
Median TTFT (ms):                        634.65    
P99 TTFT (ms):                           742.71    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          12.82     
Median TPOT (ms):                        12.71     
P99 TPOT (ms):                           13.56     
---------------Inter-token Latency----------------
Mean ITL (ms):                           12.82     
Median ITL (ms):                         12.19     
P99 ITL (ms):                            14.24     
==================================================

after

============ Serving Benchmark Result ============
Successful requests:                     200       
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  198.17    
Total input tokens:                      316188    
Total generated tokens:                  120000    
Request throughput (req/s):              1.01      
Output token throughput (tok/s):         605.55    
Peak output token throughput (tok/s):    704.00    
Peak concurrent requests:                16.00     
Total token throughput (tok/s):          2201.13   
---------------Time to First Token----------------
Mean TTFT (ms):                          677.33    
Median TTFT (ms):                        712.74    
P99 TTFT (ms):                           786.21    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          12.10     
Median TPOT (ms):                        12.02     
P99 TPOT (ms):                           12.85     
---------------Inter-token Latency----------------
Mean ITL (ms):                           12.10     
Median ITL (ms):                         11.91     
P99 ITL (ms):                            13.09     
==================================================

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request enables tunable control for naive_block_assignment in the LoRA Fused MoE layer, allowing the strategy to be specified directly in the operator configurations. While the logic for falling back to the heuristic is sound, there is a critical bug in how the configuration keys are accessed. Because the configuration dictionaries are normalized to uppercase keys before use, the current code fails to find the lowercase "naive_block_assignment" key, effectively rendering the tunable control feature non-functional. I have provided suggestions to fix this key mismatch in both the activation and summation decorators.

[gemini-code-assist]

The keys in shrink_config and expand_config are normalized to uppercase by _normalize_keys (called within _get_lora_moe_configs at line 188). Consequently, popping the lowercase "naive_block_assignment" will always return None, causing the code to always fall back to the heuristic and ignoring any tuned values. You should use the normalized key "NAIVE_BLOCK_ASSIGNMENT" instead.

Suggested change

	shrink_naive = shrink_config.pop("naive_block_assignment", None)
	expand_naive = expand_config.pop("naive_block_assignment", None)
	shrink_naive = shrink_config.pop("NAIVE_BLOCK_ASSIGNMENT", None)
	expand_naive = expand_config.pop("NAIVE_BLOCK_ASSIGNMENT", None)

[gemini-code-assist]

Similar to the issue in the activation decorator, the keys here are normalized to uppercase. Failing to pop the correct key means "NAIVE_BLOCK_ASSIGNMENT" remains in the configuration dictionaries, which may cause issues if passed down to underlying functions that do not expect it. Use the uppercase key name here as well.

Suggested change

	shrink_config.pop("naive_block_assignment", None)
	expand_config.pop("naive_block_assignment", None)
	shrink_config.pop("NAIVE_BLOCK_ASSIGNMENT", None)
	expand_config.pop("NAIVE_BLOCK_ASSIGNMENT", None)

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[github-actions]

This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you!