[gcanlin]

Purpose

Before:

  Routed MoE ──> routed_partial ─┐
                                 ├─> torch.add_ ─> Lamport RS + SConv + AG + Norm
  Shared Sink ─> shared_partial ─┘

After:

  Routed MoE ──> routed_partial ─┐
                                 ├─> Lamport RS + SConv + AG + Norm
  Shared Sink ─> shared_partial ─┘      └─ fused add during publish

Test Plan

• Benchmark for kernel only(The code below is generated by AI for testing)
• E2E

VLLM_USE_V2_MODEL_RUNNER=1 \
  FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1 \
  vllm serve thinkingmachines/Inkling-Small \
    --served-model-name thinkingmachines/Inkling-Small \
    --tokenizer-mode inkling \
    --reasoning-parser inkling \
    --tool-call-parser inkling \
    --enable-auto-tool-choice \
    --tensor-parallel-size 8 \
    --gpu-memory-utilization 0.90 \
    --kernel-config.enable_flashinfer_autotune=False \
    --trust-remote-code \
    --host 0.0.0.0 \
    --port 8000

vllm bench serve \
    --backend vllm \
    --base-url http://127.0.0.1:8000 \
    --endpoint /v1/completions \
    --model thinkingmachines/Inkling-Small \
    --tokenizer /inspire/hdd/project/embodied-multimodality/public/canlin/models/Inkling-Small \
    --dataset-name random \
    --random-input-len 32 \
    --random-output-len 1024 \
    --num-prompts 32 \
    --request-rate 1 \
    --max-concurrency 1 \
    --ignore-eos \
    --percentile-metrics ttft,tpot,itl,e2el \
    --save-result

Test Result

Kernel:

{
    "gpu": "NVIDIA H200",
    "world_size": 8,
    "hidden_size": 4096,
    "graph_repeats": 200,
    "samples": 50,
    "moe_layers": 40,
    "correctness": "exact",
    "results": [
      {
        "tokens": 1,
        "baseline_us": 8.543760180473328,
        "fused_us": 6.939600110054016,
        "saved_us_per_moe_layer": 1.6041600704193115,
        "reduction_percent": 18.775808736832317,
        "speedup": 1.2311603039050076,
        "estimated_saved_ms_per_forward": 0.06416640281677247
      }
    ]
  }

E2E:

This PR:

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             1         
Request rate configured (RPS):           1.00      
Benchmark duration (s):                  165.70    
Total input tokens:                      1024      
Total generated tokens:                  32768     
Request throughput (req/s):              0.19      
Output token throughput (tok/s):         197.75    
Peak output token throughput (tok/s):    236.00    
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          203.93    
---------------Time to First Token----------------
Mean TTFT (ms):                          16.35     
Median TTFT (ms):                        16.17     
P99 TTFT (ms):                           19.63     
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          5.03      
Median TPOT (ms):                        5.02      
P99 TPOT (ms):                           5.09      
---------------Inter-token Latency----------------
Mean ITL (ms):                           5.03      
Median ITL (ms):                         5.37      
P99 ITL (ms):                            5.64      
----------------End-to-end Latency----------------
Mean E2EL (ms):                          5156.98   
Median E2EL (ms):                        5151.16   
P99 E2EL (ms):                           5222.24   
==================================================

Main:

============ Serving Benchmark Result ============
Successful requests:                     32        
Failed requests:                         0         
Maximum request concurrency:             1         
Request rate configured (RPS):           1.00      
Benchmark duration (s):                  166.96    
Total input tokens:                      1024      
Total generated tokens:                  32768     
Request throughput (req/s):              0.19      
Output token throughput (tok/s):         196.27    
Peak output token throughput (tok/s):    234.00    
Peak concurrent requests:                2.00      
Total token throughput (tok/s):          202.40    
---------------Time to First Token----------------
Mean TTFT (ms):                          16.17     
Median TTFT (ms):                        16.02     
P99 TTFT (ms):                           18.98     
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          5.06      
Median TPOT (ms):                        5.06      
P99 TPOT (ms):                           5.07      
---------------Inter-token Latency----------------
Mean ITL (ms):                           5.12      
Median ITL (ms):                         5.40      
P99 ITL (ms):                            5.68      
----------------End-to-end Latency----------------
Mean E2EL (ms):                          5196.24   
Median E2EL (ms):                        5195.76   
P99 E2EL (ms):                           5203.33   
==================================================

---

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.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gcanlin]

@WoosukKwon @yewentao256 Would you mind taking a look? A small optimization :)

[Isotr0py]

LGTM