Batch-invariance generation with vLLM fused-MoE backend at improved performance - #6521

Merged

wdykas merged 26 commits into

NVIDIA:mainfrom

utkarsh530:batch-invariant-inference-ep8tp1

Aug 22, 2026

Commits on Aug 20, 2026

Batch-invariant mode: 64-multiple floor for CUDA-graph token buckets

Utkarsh Utkarsh
and
claude
committed
04bf0c3
Batch-invariant MoE: support gated SwiGLU activations

Utkarsh Utkarsh
and
claude
committed
b6479a7
Batch-invariant mode: support the vLLM Triton fused-MoE backend

Utkarsh Utkarsh
and
claude
committed
993cdec
Batch-invariant mode: add 'te_native' GEMM backend (workspace starvation)

Utkarsh Utkarsh
and
claude
committed
3b62e8d
Batch-invariant mode: make the GEMM backend selectable via config

Utkarsh Utkarsh
and
claude
committed
6b9e1d6
Batch-invariant te_native backend: keep native TE RMSNorm

Utkarsh Utkarsh
and
claude
committed
da6ad41
Batch-invariant vLLM MoE backend: pin only the K-reduction recipe

Utkarsh Utkarsh
and
claude
committed
c4166a5
Tests: batch-invariant vLLM fused-MoE backend

Utkarsh Utkarsh
and
claude
committed
3e85a90
Docstrings: document te_native backend and _moe_sum options

Utkarsh Utkarsh
and
claude
committed
1ed83b9
Apply tools/autoformat.sh (black/isort/ruff)

Utkarsh Utkarsh
and
claude
committed
d5ece84
Address review: clearer rounding expression in _batch_invariant_token_floor

Utkarsh Utkarsh
and
claude
committed
ad1c023
Address review: make te_native the default batch-invariant backend

Utkarsh Utkarsh
and
claude
committed
cc58495
Address review: make _moe_sum apply_weights/acc_fp64 truly orthogonal

Utkarsh Utkarsh
and
claude
committed
41f8dd3
Address review: pad buckets to alignment instead of invalidating them

Utkarsh Utkarsh
and
claude
committed
e5bc58b
Address review: restore te_native workspace state on disable

Utkarsh Utkarsh
and
claude
committed
93e933c
Address review: scope acc_fp64 to SwiGLU; type the backend field

Utkarsh Utkarsh
and
claude
committed
ed496e8
Address review: derive weighted-swiglu grid size from SM count
Utkarsh Utkarsh
committed
ed183b7
Address review: validate batch_invariant_backend in __post_init__
Utkarsh Utkarsh
committed
da0407b
Address review: finish align-up wording; tie rounder to TOKEN_ROUNDER
Utkarsh Utkarsh
committed
371772d
Address review: scope skip_rmsnorm to normalization patches only
Utkarsh Utkarsh
committed
e69fe65
Address review: skip BI kernel tests without CUDA/Triton
Utkarsh Utkarsh
committed
2cdc4b6
Address review: drop getattr default for batch_invariant_backend
Utkarsh Utkarsh
committed
bf46b6e
Address review: _moe_sum docstring — fp32 or fp64 accumulation
Utkarsh Utkarsh
committed
e3aa438
Address review: single module-level TOKEN_ROUNDER constant
Utkarsh Utkarsh
committed
9051751
Apply tools/autoformat.sh (isort import order in BI test)
Utkarsh Utkarsh
committed
8a25689

Commits on Aug 21, 2026

Add batch_invariant_backend to Mamba MoE golden config
Utkarsh Utkarsh
committed
a5aa088

Navigation Menu

Uh oh!

Batch-invariance generation with vLLM fused-MoE backend at improved performance - #6521

Batch-invariance generation with vLLM fused-MoE backend at improved performance#6521
wdykas merged 26 commits into
NVIDIA:mainfrom
utkarsh530:batch-invariant-inference-ep8tp1

Commits on Aug 20, 2026

Batch-invariant mode: 64-multiple floor for CUDA-graph token buckets

Batch-invariant MoE: support gated SwiGLU activations

Batch-invariant mode: support the vLLM Triton fused-MoE backend

Batch-invariant mode: add 'te_native' GEMM backend (workspace starvation)

Batch-invariant mode: make the GEMM backend selectable via config

Batch-invariant te_native backend: keep native TE RMSNorm

Batch-invariant vLLM MoE backend: pin only the K-reduction recipe

Tests: batch-invariant vLLM fused-MoE backend

Docstrings: document te_native backend and _moe_sum options

Apply tools/autoformat.sh (black/isort/ruff)

Address review: clearer rounding expression in _batch_invariant_token_floor

Address review: make te_native the default batch-invariant backend

Address review: make _moe_sum apply_weights/acc_fp64 truly orthogonal

Address review: pad buckets to alignment instead of invalidating them

Address review: restore te_native workspace state on disable

Address review: scope acc_fp64 to SwiGLU; type the backend field

Address review: derive weighted-swiglu grid size from SM count

Address review: validate batch_invariant_backend in __post_init__

Address review: finish align-up wording; tie rounder to TOKEN_ROUNDER

Address review: scope skip_rmsnorm to normalization patches only

Address review: skip BI kernel tests without CUDA/Triton

Address review: drop getattr default for batch_invariant_backend

Address review: _moe_sum docstring — fp32 or fp64 accumulation

Address review: single module-level TOKEN_ROUNDER constant

Apply tools/autoformat.sh (isort import order in BI test)

Commits on Aug 21, 2026

Add batch_invariant_backend to Mamba MoE golden config