[Metrics] Add Prometheus counter for CUDA graph iteration mode

Add `vllm:cudagraph_iterations` counter with `runtime_mode` label (NONE/PIECEWISE/FULL) so CUDA graph usage percentage can be computed in PromQL. Also always populate CUDAGraphStat regardless of the `--cudagraph-metrics` flag, since the stat is trivially cheap and that flag should only gate the verbose text log table. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
[bug] Fix deadlock with pause resume and collective_rpc (#37024 )
2026-03-18 23:11:25 -04:00 · 2026-03-19 01:49:02 +00:00 · 2026-03-19 01:36:27 +00:00 · 2026-03-19 00:45:10 +00:00 · 2026-03-18 20:34:00 -04:00
35 changed files with 773 additions and 1786 deletions
@@ -1,116 +0,0 @@
-# NVFP4 NaN Contamination Fix
-
-## Summary
-
-Fixed a critical bug in NVFP4 quantization where NaN values in input tensors caused 100% of the output to become NaN.
-
-## The Bug
-
-**Root Cause**: When a tensor contains NaN in any block (e.g., from attention softmax producing 0/0), the FP4 block scale for that block becomes NaN. During the GEMM operation, this NaN block scale contaminates the **entire output** for that token.
-
-**Reproduction**:
-```python
-# Input: Single token with NaN in block 1 (dims 16-31)
-x = torch.randn(1, 64, dtype=torch.bfloat16)
-x[0, 16:32] = float('nan')
-
-# After quantization:
-# Block 0 scale: 0.375 (clean)
-# Block 1 scale: NaN   ← Problem!
-# Block 2 scale: 0.281 (clean)
-# Block 3 scale: 0.219 (clean)
-
-# After GEMM: 100% of output is NaN
-```
-
-## The Fix
-
-**Location**: `vllm/model_executor/layers/quantization/utils/nvfp4_utils.py:219`
-
-**Change**: Added NaN masking before FP4 quantization:
-
-```python
-# Mask NaNs before quantization to prevent block scale contamination
-x = torch.where(torch.isnan(x), torch.zeros_like(x), x)
-```
-
-**Why it works**:
- NaN → 0 prevents NaN from contaminating block scales
- Zero-cost operation (compiles to a single select instruction)
- Preserves clean data while safely handling NaN inputs
-
-## Test Coverage
-
-### 1. **test_nvfp4_nan_block_contamination.py** - Demonstrates the bug
- ❌ **Buggy path** (`use_fix=False`): 100% of output is NaN
- ✅ **Fixed path** (`use_fix=True`): 0% of output is NaN
-
-### 2. **test_nvfp4_nan_integration.py** - Integration test
- ✅ Verifies production code fix through full `apply_nvfp4_linear()` path
- Input with NaN → Clean output (no NaN contamination)
-
-### 3. **test_nvfp4_nan_propagation.py** - Comprehensive test suite
- Tests multiple NaN placement strategies (end, middle, scattered)
- Tests various batch sizes, hidden dims, and data types
- Validates both buggy and fixed code paths
-
-## Results
-
-**Before fix**:
-```
-Block 1 scale: nan
-Output: [nan, nan, nan, nan, ..., nan]  (100% NaN)
-```
-
-**After fix**:
-```
-Block 1 scale: 0.0
-Output: [3014656., -4587520., -1515520., ...]  (0% NaN)
-```
-
-## Regression Testing
-
-All existing NVFP4 tests pass:
- ✅ `test_nvfp4_quant.py`: 50/50 tests passed
- ✅ `test_nvfp4_scaled_mm.py`: 12/12 tests passed
- ✅ No performance impact (zero-cost NaN masking)
-
-## Impact
-
- **Fixes**: Wide EP DeepSeek R1 NaN crashes on GB200s
- **Prevents**: Future NaN contamination from attention/softmax operations
- **Cost**: ~19us per layer (~0.6ms for 32-layer model)
-  - Overhead: ~50% on the quantization step itself
-  - Negligible in practice: 0.6ms vs model crashing with 100% NaN
-  - Cannot fuse into custom CUDA op without kernel changes
- **Fullgraph compatible**: Simple element-wise operation, no graph breaks
-
-## Future Optimization
-
-If the ~19us/layer overhead becomes significant, we can:
-1. **Integrate into CUDA kernel**: Modify `scaled_fp4_quant` to mask NaN during load (true zero-cost)
-2. **Integrate with check_tensor**: Add `replace_nan=True` parameter to existing NaN detector
-3. **Upstream masking**: Fix attention layer to never produce NaN in the first place
-
-For now, the trade-off is acceptable: ~0.6ms overhead vs 100% NaN crash.
-
-## Files Changed
-
-1. **vllm/model_executor/layers/quantization/utils/nvfp4_utils.py**
-   - Added NaN masking in `apply_nvfp4_linear()` before quantization
-
-2. **tests/kernels/quantization/test_nvfp4_nan_block_contamination.py** (new)
-   - Demonstrates the bug and validates the fix
-
-3. **tests/kernels/quantization/test_nvfp4_nan_integration.py** (new)
-   - End-to-end integration test through production code path
-
-4. **tests/kernels/quantization/test_nvfp4_nan_propagation.py** (new)
-   - Comprehensive test suite for various NaN scenarios
-
---
-
-**Date**: 2026-03-28
-**Author**: Claude Sonnet 4.5
-**Issue**: NaN contamination in NVFP4 o_proj GEMM
-**Status**: Fixed and tested ✅
@@ -20,8 +20,7 @@ __global__ void rms_norm_kernel(
    const int64_t input_shape_d2,         // input.size(-2)
    const int64_t input_shape_d3,         // input.size(-3)
    const scalar_t* __restrict__ weight,  // [hidden_size]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  __shared__ float s_variance;
  float variance = 0.0f;
  const scalar_t* input_row;
@@ -64,9 +63,6 @@ __global__ void rms_norm_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -98,8 +94,7 @@ fused_add_rms_norm_kernel(
    const int64_t input_stride,
    scalar_t* __restrict__ residual,      // [..., hidden_size]
    const scalar_t* __restrict__ weight,  // [hidden_size]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  // Sanity checks on our vector struct and type-punned pointer arithmetic
  static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
  static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);
@@ -133,9 +128,6 @@ fused_add_rms_norm_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -159,8 +151,7 @@ fused_add_rms_norm_kernel(
    const int64_t input_stride,
    scalar_t* __restrict__ residual,      // [..., hidden_size]
    const scalar_t* __restrict__ weight,  // [hidden_size]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  __shared__ float s_variance;
  float variance = 0.0f;

@@ -178,9 +169,6 @@ fused_add_rms_norm_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -196,10 +184,7 @@ fused_add_rms_norm_kernel(
 void rms_norm(torch::Tensor& out,     // [..., hidden_size]
              torch::Tensor& input,   // [..., hidden_size]
              torch::Tensor& weight,  // [hidden_size]
-              double epsilon,
-              std::optional<torch::Tensor> nan_flags,
-              int64_t layer_idx,
-              int64_t max_num_tokens) {
+              double epsilon) {
  TORCH_CHECK(out.is_contiguous());
  if (input.stride(-1) != 1) {
    input = input.contiguous();
@@ -217,11 +202,6 @@ void rms_norm(torch::Tensor& out,     // [..., hidden_size]
  int64_t input_shape_d2 = (num_dims >= 3) ? input.size(-2) : 0;
  int64_t input_shape_d3 = (num_dims >= 4) ? input.size(-3) : 0;

-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr = nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  // For large num_tokens, use smaller blocks to increase SM concurrency.
  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
  dim3 grid(num_tokens);
@@ -240,7 +220,7 @@ void rms_norm(torch::Tensor& out,     // [..., hidden_size]
                out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
                input_stride_d2, input_stride_d3, input_stride_d4,
                input_shape_d2, input_shape_d3, weight.data_ptr<scalar_t>(),
-                epsilon, num_tokens, hidden_size, nan_flag_ptr);
+                epsilon, num_tokens, hidden_size);
      });
    });
  });
@@ -253,16 +233,13 @@ void rms_norm(torch::Tensor& out,     // [..., hidden_size]
            <<<grid, block, 0, stream>>>(                                   \
                input.data_ptr<scalar_t>(), input_stride,                   \
                residual.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(), \
-                epsilon, num_tokens, hidden_size, nan_flag_ptr);            \
+                epsilon, num_tokens, hidden_size);                          \
      });

 void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
                        torch::Tensor& residual,  // [..., hidden_size]
                        torch::Tensor& weight,    // [hidden_size]
-                        double epsilon,
-                        std::optional<torch::Tensor> nan_flags,
-                        int64_t layer_idx,
-                        int64_t max_num_tokens) {
+                        double epsilon) {
  TORCH_CHECK(weight.scalar_type() == input.scalar_type());
  TORCH_CHECK(input.scalar_type() == residual.scalar_type());
  TORCH_CHECK(residual.is_contiguous());
@@ -271,11 +248,6 @@ void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
  int64_t input_stride = input.stride(-2);
  int num_tokens = input.numel() / hidden_size;

-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr = nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  dim3 grid(num_tokens);
  /* This kernel is memory-latency bound in many scenarios.
     When num_tokens is large, a smaller block size allows
@@ -25,8 +25,7 @@ __global__ void rms_norm_static_fp8_quant_kernel(
    const int input_stride,
    const scalar_t* __restrict__ weight,  // [hidden_size]
    const float* __restrict__ scale,      // [1]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  __shared__ float s_variance;
  float variance = 0.0f;

@@ -52,9 +51,6 @@ __global__ void rms_norm_static_fp8_quant_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -89,8 +85,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
    scalar_t* __restrict__ residual,      // [..., hidden_size]
    const scalar_t* __restrict__ weight,  // [hidden_size]
    const float* __restrict__ scale,      // [1]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  // Sanity checks on our vector struct and type-punned pointer arithmetic
  static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
  static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);
@@ -124,9 +119,6 @@ fused_add_rms_norm_static_fp8_quant_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -158,8 +150,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
    scalar_t* __restrict__ residual,      // [..., hidden_size]
    const scalar_t* __restrict__ weight,  // [hidden_size]
    const float* __restrict__ scale,      // [1]
-    const float epsilon, const int num_tokens, const int hidden_size,
-    int8_t* __restrict__ nan_flag_ptr) {
+    const float epsilon, const int num_tokens, const int hidden_size) {
  __shared__ float s_variance;
  float variance = 0.0f;

@@ -177,9 +168,6 @@ fused_add_rms_norm_static_fp8_quant_kernel(

  if (threadIdx.x == 0) {
    s_variance = rsqrtf(variance / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(variance) || isinf(variance))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -200,20 +188,12 @@ void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
                               torch::Tensor& input,   // [..., hidden_size]
                               torch::Tensor& weight,  // [hidden_size]
                               torch::Tensor& scale,   // [1]
-                               double epsilon,
-                               std::optional<torch::Tensor> nan_flags,
-                               int64_t layer_idx,
-                               int64_t max_num_tokens) {
+                               double epsilon) {
  TORCH_CHECK(out.is_contiguous());
  int hidden_size = input.size(-1);
  int input_stride = input.stride(-2);
  int num_tokens = input.numel() / hidden_size;

-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr = nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  // For large num_tokens, use smaller blocks to increase SM concurrency.
  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
  dim3 grid(num_tokens);
@@ -235,7 +215,7 @@ void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
                        out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
                        input_stride, weight.data_ptr<scalar_t>(),
                        scale.data_ptr<float>(), epsilon, num_tokens,
-                        hidden_size, nan_flag_ptr);
+                        hidden_size);
              });
            });
      });
@@ -252,7 +232,7 @@ void rms_norm_static_fp8_quant(torch::Tensor& out,     // [..., hidden_size]
                      out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),     \
                      input_stride, residual.data_ptr<scalar_t>(),           \
                      weight.data_ptr<scalar_t>(), scale.data_ptr<float>(),  \
-                      epsilon, num_tokens, hidden_size, nan_flag_ptr);       \
+                      epsilon, num_tokens, hidden_size);                     \
            });                                                              \
      });
 void fused_add_rms_norm_static_fp8_quant(
@@ -261,10 +241,7 @@ void fused_add_rms_norm_static_fp8_quant(
    torch::Tensor& residual,  // [..., hidden_size]
    torch::Tensor& weight,    // [hidden_size]
    torch::Tensor& scale,     // [1]
-    double epsilon,
-    std::optional<torch::Tensor> nan_flags,
-    int64_t layer_idx,
-    int64_t max_num_tokens) {
+    double epsilon) {
  TORCH_CHECK(out.is_contiguous());
  TORCH_CHECK(residual.is_contiguous());
  TORCH_CHECK(residual.scalar_type() == input.scalar_type());
@@ -273,11 +250,6 @@ void fused_add_rms_norm_static_fp8_quant(
  int input_stride = input.stride(-2);
  int num_tokens = input.numel() / hidden_size;

-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr = nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  dim3 grid(num_tokens);
  /* This kernel is memory-latency bound in many scenarios.
     When num_tokens is large, a smaller block size allows
@@ -87,14 +87,10 @@ void convert_vertical_slash_indexes_mergehead(
 #endif

 void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
-              double epsilon,
-              std::optional<torch::Tensor> nan_flags = std::nullopt,
-              int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+              double epsilon);

 void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
-                        torch::Tensor& weight, double epsilon,
-                        std::optional<torch::Tensor> nan_flags = std::nullopt,
-                        int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+                        torch::Tensor& weight, double epsilon);

 void fused_qk_norm_rope(torch::Tensor& qkv, int64_t num_heads_q,
                        int64_t num_heads_k, int64_t num_heads_v,
@@ -124,17 +120,13 @@ void large_context_topk(const torch::Tensor& score, torch::Tensor& indices,

 void rms_norm_static_fp8_quant(torch::Tensor& out, torch::Tensor& input,
                               torch::Tensor& weight, torch::Tensor& scale,
-                               double epsilon,
-                               std::optional<torch::Tensor> nan_flags = std::nullopt,
-                               int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+                               double epsilon);

 void fused_add_rms_norm_static_fp8_quant(torch::Tensor& out,
                                         torch::Tensor& input,
                                         torch::Tensor& residual,
                                         torch::Tensor& weight,
-                                         torch::Tensor& scale, double epsilon,
-                                         std::optional<torch::Tensor> nan_flags = std::nullopt,
-                                         int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+                                         torch::Tensor& scale, double epsilon);

 void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
                                      torch::Tensor const& input,
@@ -142,18 +134,14 @@ void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
                                      torch::Tensor& scales,
                                      double const epsilon,
                                      std::optional<torch::Tensor> scale_ub,
-                                      std::optional<torch::Tensor> residual,
-                                      std::optional<torch::Tensor> nan_flags = std::nullopt,
-                                      int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+                                      std::optional<torch::Tensor> residual);

 void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
                              torch::Tensor const& weight,
                              torch::Tensor& scales, double const epsilon,
                              std::optional<torch::Tensor> scale_ub,
                              std::optional<torch::Tensor> residual,
-                              int64_t group_size, bool is_scale_transposed,
-                              std::optional<torch::Tensor> nan_flags = std::nullopt,
-                              int64_t layer_idx = 0, int64_t max_num_tokens = 0);
+                              int64_t group_size, bool is_scale_transposed);

 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
                      std::optional<torch::Tensor> key, int64_t head_size,
@@ -15,15 +15,13 @@ __device__ void rms_norm_dynamic_per_token_quant_vec(
    scalar_t const* __restrict__ input,   // [..., hidden_size]
    scalar_t const* __restrict__ weight,  // [hidden_size]
    float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
-    int32_t const input_stride, scalar_t* __restrict__ residual = nullptr,
-    int8_t* __restrict__ nan_flag_ptr = nullptr) {
+    int32_t const input_stride, scalar_t* __restrict__ residual = nullptr) {
  float rms = 0.0f;
  float token_scale = 0.0f;

  // Compute rms
  vllm::vectorized::compute_rms<scalar_t, has_residual>(
-      &rms, input, hidden_size, input_stride, var_epsilon, residual,
-      nan_flag_ptr);
+      &rms, input, hidden_size, input_stride, var_epsilon, residual);

  // Compute scale
  vllm::vectorized::compute_dynamic_per_token_scales<scalar_t, scalar_out_t,
@@ -55,8 +53,7 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
    scalar_t const* __restrict__ input,   // [..., hidden_size]
    scalar_t const* __restrict__ weight,  // [hidden_size]
    float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
-    int32_t const input_stride, scalar_t* __restrict__ residual = nullptr,
-    int8_t* __restrict__ nan_flag_ptr = nullptr) {
+    int32_t const input_stride, scalar_t* __restrict__ residual = nullptr) {
  // For vectorization, token_input and token_output pointers need to be
  // aligned at 8-byte and 4-byte addresses respectively.
  bool const can_vectorize = hidden_size % 4 == 0 and input_stride % 4 == 0;
@@ -65,7 +62,7 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
    return rms_norm_dynamic_per_token_quant_vec<scalar_t, scalar_out_t,
                                                has_residual>(
        out, scales, input, weight, scale_ub, var_epsilon, hidden_size,
-        input_stride, residual, nan_flag_ptr);
+        input_stride, residual);
  }

  float rms = 0.0f;
@@ -73,8 +70,7 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(

  // Compute RMS
  vllm::compute_rms<scalar_t, has_residual>(
-      &rms, input, hidden_size, input_stride, var_epsilon, residual,
-      nan_flag_ptr);
+      &rms, input, hidden_size, input_stride, var_epsilon, residual);
  // Compute Scale
  vllm::compute_dynamic_per_token_scales<scalar_t, scalar_out_t, has_residual>(
      &token_scale, scales, input, weight, rms, scale_ub, hidden_size,
@@ -106,14 +102,12 @@ __global__ void rms_norm_per_block_quant_kernel(
    scalar_t const* __restrict__ weight,  // [hidden_size]
    float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
    int32_t const input_stride, scalar_t* __restrict__ residual = nullptr,
-    int64_t outer_scale_stride = 1,
-    int8_t* __restrict__ nan_flag_ptr = nullptr) {
+    int64_t outer_scale_stride = 1) {
  float rms;
  // Compute RMS
  // Always able to vectorize due to constraints on hidden_size
  vllm::vectorized::compute_rms<scalar_t, has_residual>(
-      &rms, input, hidden_size, input_stride, var_epsilon, residual,
-      nan_flag_ptr);
+      &rms, input, hidden_size, input_stride, var_epsilon, residual);

  // Compute Scale
  // Always able to vectorize due to constraints on hidden_size and group_size
@@ -146,8 +140,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
    torch::Tensor& scales,        // [num_tokens]
    double const var_epsilon,     // Variance epsilon used in norm calculation
    std::optional<at::Tensor> const& scale_ub,
-    std::optional<at::Tensor>& residual,
-    int8_t* nan_flag_ptr) {
+    std::optional<at::Tensor>& residual) {
  int32_t hidden_size = input.size(-1);
  int32_t input_stride = input.view({-1, hidden_size}).stride(0);
  auto num_tokens = input.numel() / hidden_size;
@@ -167,8 +160,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
                  input.data_ptr<scalar_in_t>(), weight.data_ptr<scalar_in_t>(),
                  scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
                  var_epsilon, hidden_size, input_stride,
-                  has_residual ? residual->data_ptr<scalar_in_t>() : nullptr,
-                  nan_flag_ptr);
+                  has_residual ? residual->data_ptr<scalar_in_t>() : nullptr);
        });
  });
 }
@@ -179,9 +171,7 @@ void rms_norm_dynamic_per_token_quant(
    torch::Tensor const& weight,  // [hidden_size]
    torch::Tensor& scales,        // [num_tokens]
    double const var_epsilon,     // Variance epsilon used in norm calculation
-    std::optional<at::Tensor> scale_ub, std::optional<at::Tensor> residual,
-    std::optional<torch::Tensor> nan_flags, int64_t layer_idx,
-    int64_t max_num_tokens) {
+    std::optional<at::Tensor> scale_ub, std::optional<at::Tensor> residual) {
  static c10::ScalarType kFp8Type = is_fp8_ocp()
                                        ? c10::ScalarType::Float8_e4m3fn
                                        : c10::ScalarType::Float8_e4m3fnuz;
@@ -200,17 +190,10 @@ void rms_norm_dynamic_per_token_quant(
    TORCH_CHECK(residual->is_contiguous());
  }

-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr =
-        nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  VLLM_DISPATCH_FLOATING_TYPES(
      input.scalar_type(), "rms_norm_dynamic_per_token_quant_dispatch", [&] {
        rms_norm_dynamic_per_token_quant_dispatch<scalar_t>(
-            out, input, weight, scales, var_epsilon, scale_ub, residual,
-            nan_flag_ptr);
+            out, input, weight, scales, var_epsilon, scale_ub, residual);
      });
 }

@@ -224,8 +207,7 @@ void rms_norm_per_block_quant_dispatch(
    int32_t group_size,
    double const var_epsilon,  // Variance epsilon used in norm calculation
    std::optional<at::Tensor> const& scale_ub,
-    std::optional<at::Tensor>& residual, bool is_scale_transposed,
-    int8_t* nan_flag_ptr) {
+    std::optional<at::Tensor>& residual, bool is_scale_transposed) {
  int32_t hidden_size = input.size(-1);
  int32_t input_stride = input.view({-1, hidden_size}).stride(0);

@@ -264,7 +246,7 @@ void rms_norm_per_block_quant_dispatch(
                            var_epsilon, hidden_size, input_stride,
                            has_residual ? residual->data_ptr<scalar_in_t>()
                                         : nullptr,
-                            scales.stride(1), nan_flag_ptr);
+                            scales.stride(1));
                  });
            });
          });
@@ -277,9 +259,7 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
                              torch::Tensor& scales, double const var_epsilon,
                              std::optional<torch::Tensor> scale_ub,
                              std::optional<torch::Tensor> residual,
-                              int64_t group_size, bool is_scale_transposed,
-                              std::optional<torch::Tensor> nan_flags,
-                              int64_t layer_idx, int64_t max_num_tokens) {
+                              int64_t group_size, bool is_scale_transposed) {
  static c10::ScalarType kFp8Type = is_fp8_ocp()
                                        ? c10::ScalarType::Float8_e4m3fn
                                        : c10::ScalarType::Float8_e4m3fnuz;
@@ -315,13 +295,7 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
              "scales buffer too small: need ", num_tokens * num_groups,
              " elements, got ", scales.numel());

-  int8_t* nan_flag_ptr = nullptr;
-  if (nan_flags.has_value()) {
-    nan_flag_ptr =
-        nan_flags->data_ptr<int8_t>() + layer_idx * max_num_tokens;
-  }
-
  rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
                                    var_epsilon, scale_ub, residual,
-                                    is_scale_transposed, nan_flag_ptr);
+                                    is_scale_transposed);
 }
@@ -18,8 +18,7 @@ template <typename scalar_t, bool has_residual = false>
 __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
                            int32_t const hidden_size,
                            int32_t const input_stride, float const epsilon,
-                            scalar_t const* __restrict__ residual = nullptr,
-                            int8_t* __restrict__ nan_flag_ptr = nullptr) {
+                            scalar_t const* __restrict__ residual = nullptr) {
  int64_t const input_token_offset =
      blockIdx.x * static_cast<int64_t>(input_stride);
  int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
@@ -42,9 +41,6 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
  __shared__ float s_rms;
  if (threadIdx.x == 0) {
    s_rms = rsqrtf(ss / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(ss) || isinf(ss))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -239,8 +235,7 @@ template <typename scalar_t, bool has_residual = false>
 __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
                            int32_t const hidden_size,
                            int32_t const input_stride, float const epsilon,
-                            scalar_t const* __restrict__ residual = nullptr,
-                            int8_t* __restrict__ nan_flag_ptr = nullptr) {
+                            scalar_t const* __restrict__ residual = nullptr) {
  int64_t const input_token_offset =
      blockIdx.x * static_cast<int64_t>(input_stride);
  int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
@@ -291,9 +286,6 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
  __shared__ float s_rms;
  if (threadIdx.x == 0) {
    s_rms = rsqrtf(ss / hidden_size + epsilon);
-    if (nan_flag_ptr && (isnan(ss) || isinf(ss))) {
-      nan_flag_ptr[blockIdx.x] = 1;
-    }
  }
  __syncthreads();

@@ -152,15 +152,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // Layernorm
  // Apply Root Mean Square (RMS) Normalization to the input tensor.
  ops.def(
-      "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon, "
-      "Tensor? nan_flags=None, int layer_idx=0, int max_num_tokens=0) -> ()");
+      "rms_norm(Tensor! result, Tensor input, Tensor weight, float epsilon) -> "
+      "()");
  ops.impl("rms_norm", torch::kCUDA, &rms_norm);

  // In-place fused Add and RMS Normalization.
  ops.def(
      "fused_add_rms_norm(Tensor! input, Tensor! residual, Tensor weight, "
-      "float epsilon, Tensor? nan_flags=None, int layer_idx=0, "
-      "int max_num_tokens=0) -> ()");
+      "float epsilon) -> ()");
  ops.impl("fused_add_rms_norm", torch::kCUDA, &fused_add_rms_norm);

  // Function for fused QK Norm and RoPE
@@ -201,8 +200,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  // Apply Root Mean Square (RMS) Normalization to the input tensor.
  ops.def(
      "rms_norm_static_fp8_quant(Tensor! result, Tensor input, Tensor weight, "
-      "Tensor scale, float epsilon, Tensor? nan_flags=None, "
-      "int layer_idx=0, int max_num_tokens=0) -> ()");
+      "Tensor scale, float epsilon) -> "
+      "()");
  ops.impl("rms_norm_static_fp8_quant", torch::kCUDA,
           &rms_norm_static_fp8_quant);

@@ -210,8 +209,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "fused_add_rms_norm_static_fp8_quant(Tensor! result, Tensor input, "
      "Tensor! residual, Tensor weight, "
-      "Tensor scale, float epsilon, Tensor? nan_flags=None, "
-      "int layer_idx=0, int max_num_tokens=0) -> ()");
+      "Tensor scale, float epsilon) -> ()");
  ops.impl("fused_add_rms_norm_static_fp8_quant", torch::kCUDA,
           &fused_add_rms_norm_static_fp8_quant);

@@ -219,8 +217,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
  ops.def(
      "rms_norm_dynamic_per_token_quant(Tensor! result, Tensor input, "
      "Tensor weight, Tensor! scale, float epsilon, "
-      "Tensor? scale_ub, Tensor!? residual, Tensor? nan_flags=None, "
-      "int layer_idx=0, int max_num_tokens=0) -> ()");
+      "Tensor? scale_ub, Tensor!? residual) -> ()");
  ops.impl("rms_norm_dynamic_per_token_quant", torch::kCUDA,
           &rms_norm_dynamic_per_token_quant);

@@ -229,8 +226,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
      "rms_norm_per_block_quant(Tensor! result, Tensor input, "
      "Tensor weight, Tensor! scale, float epsilon, "
      "Tensor? scale_ub, Tensor!? residual, int group_size, "
-      "bool is_scale_transposed, Tensor? nan_flags=None, "
-      "int layer_idx=0, int max_num_tokens=0) -> ()");
+      "bool is_scale_transposed) -> ()");
  ops.impl("rms_norm_per_block_quant", torch::kCUDA, &rms_norm_per_block_quant);

  // Rotary embedding
@@ -0,0 +1,339 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+RLHF with FSDP2 training (4 GPUs) and vLLM expert-parallel inference (4 GPUs).
+
+8-GPU layout:
+  Training  — 4 GPUs, PyTorch FSDP2 (fully_shard)
+  Inference — 4 GPUs, vLLM AsyncLLMEngine with expert parallelism +
+              data parallelism (TP=1, DP=4, enable_expert_parallel
+              → EP_SIZE = TP×DP = 4)
+
+FSDP workers are Ray actors that form a single FSDP2 process group.
+Rank 0 gathers full parameters via DTensor.full_tensor() and broadcasts
+them to the vLLM inference engine through the NCCL weight-transfer API.
+
+The inference engine uses AsyncLLMEngine which automatically spawns
+DP worker processes (no manual placement group needed).  Weight sync
+uses pause_generation / resume_generation.
+
+Steps:
+  1. Launch 4 FSDP training workers.
+  2. Launch AsyncLLMEngine with EP+DP (dummy weights).
+  3. Generate from prompts → gibberish (random weights).
+  4. Pause generation, transfer weights from FSDP, resume.
+  5. Generate from prompts → sensible output (synced weights).
+
+Assumes a single-node cluster with 8 GPUs.
+"""
+
+import asyncio
+import os
+import uuid
+from dataclasses import asdict
+
+import ray
+import torch
+import torch.distributed as dist
+from huggingface_hub import snapshot_download
+from torch.distributed.fsdp import fully_shard
+from transformers import AutoModelForCausalLM
+
+import vllm
+from vllm import SamplingParams
+from vllm.config import WeightTransferConfig
+from vllm.distributed.weight_transfer.base import (
+    WeightTransferInitRequest,
+    WeightTransferUpdateRequest,
+)
+from vllm.distributed.weight_transfer.nccl_engine import (
+    NCCLTrainerSendWeightsArgs,
+    NCCLWeightTransferEngine,
+    NCCLWeightTransferInitInfo,
+    NCCLWeightTransferUpdateInfo,
+)
+from vllm.utils.network_utils import get_ip, get_open_port
+from vllm.v1.executor import Executor
+
+MODEL_NAME = "Qwen/Qwen3-30B-A3B"
+
+FSDP_WORLD_SIZE = 4
+INFERENCE_TP_SIZE = 1
+INFERENCE_DP_SIZE = 4
+
+
+@ray.remote(num_gpus=1)
+class FSDPTrainWorker:
+    """
+    One FSDP2 training worker per GPU.  Four of these form the FSDP group.
+    Rank 0 additionally handles weight transfer to the vLLM engine.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        rank: int,
+        fsdp_world_size: int,
+        fsdp_master_addr: str,
+        fsdp_master_port: int,
+    ):
+        self.rank = rank
+
+        os.environ["MASTER_ADDR"] = fsdp_master_addr
+        os.environ["MASTER_PORT"] = str(fsdp_master_port)
+
+        dist.init_process_group(backend="nccl", rank=rank, world_size=fsdp_world_size)
+        torch.accelerator.set_device_index(0)
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, torch_dtype=torch.bfloat16
+        )
+
+        self.weight_names = [n for n, _ in model.named_parameters()]
+        self.weight_dtype_names = [
+            str(p.dtype).split(".")[-1] for _, p in model.named_parameters()
+        ]
+        self.weight_shapes = [list(p.shape) for _, p in model.named_parameters()]
+
+        for layer in model.model.layers:
+            fully_shard(layer)
+        fully_shard(model)
+
+        self.model = model
+
+        self.transfer_port = None
+        self.transfer_master_address = None
+        self.model_update_group = None
+
+    def get_rank(self):
+        return self.rank
+
+    # ---- weight-transfer setup (rank 0 only) ----
+
+    def setup_transfer_endpoint(self):
+        """Create the NCCL rendezvous endpoint for weight transfer."""
+        assert self.rank == 0
+        self.transfer_port = get_open_port()
+        self.transfer_master_address = get_ip()
+        return self.transfer_master_address, self.transfer_port
+
+    def init_weight_transfer_group(self, transfer_world_size: int):
+        """Join the weight-transfer NCCL group as rank 0 (the source)."""
+        assert self.rank == 0
+        self.model_update_group = NCCLWeightTransferEngine.trainer_init(
+            dict(
+                master_address=self.transfer_master_address,
+                master_port=self.transfer_port,
+                world_size=transfer_world_size,
+            ),
+        )
+
+    def get_weight_metadata(self):
+        """Return weight names, dtypes, and shapes captured before FSDP wrapping."""
+        return self.weight_names, self.weight_dtype_names, self.weight_shapes
+
+    # ---- collective ops (ALL FSDP ranks must call concurrently) ----
+
+    def gather_and_broadcast_weights(self, packed: bool = True):
+        """
+        All-gather full parameters and broadcast them to vLLM.
+        Only rank 0 performs the actual NCCL broadcast; others just
+        participate in the FSDP all-gather.
+
+        full_tensor() is a collective — all FSDP ranks must call it
+        for each parameter in the same order.  Rank 0 additionally
+        feeds each gathered tensor to the weight-transfer engine.
+        """
+        if self.rank == 0:
+
+            def _full_param_iter():
+                for name, param in self.model.named_parameters():
+                    yield name, param.full_tensor()
+
+            trainer_args = NCCLTrainerSendWeightsArgs(
+                group=self.model_update_group,
+                packed=packed,
+            )
+            NCCLWeightTransferEngine.trainer_send_weights(
+                iterator=_full_param_iter(),
+                trainer_args=trainer_args,
+            )
+        else:
+            for _, param in self.model.named_parameters():
+                param.full_tensor()
+
+
+def create_async_engine(**kwargs):
+    """Create an AsyncLLMEngine directly (no subclass needed)."""
+    engine_args = vllm.AsyncEngineArgs(**kwargs)
+    vllm_config = engine_args.create_engine_config()
+    executor_class = Executor.get_class(vllm_config)
+    return vllm.AsyncLLMEngine(
+        vllm_config=vllm_config,
+        executor_class=executor_class,
+        log_requests=engine_args.enable_log_requests,
+        log_stats=not engine_args.disable_log_stats,
+    )
+
+
+async def generate_batch(engine, prompts, sampling_params):
+    """Generate completions for a batch of prompts."""
+
+    async def gen_one(prompt):
+        output = None
+        async for request_output in engine.generate(
+            {"prompt": prompt},
+            sampling_params,
+            request_id=str(uuid.uuid4()),
+        ):
+            output = request_output
+        return output
+
+    return await asyncio.gather(*[gen_one(p) for p in prompts])
+
+
+async def main():
+    ray.init()
+
+    # Download model weights to local/shared disk once.
+    local_model_path = snapshot_download(MODEL_NAME)
+    print(f"[init] Model downloaded to {local_model_path}")
+
+    # FSDP rendezvous address (single-node)
+    fsdp_master_addr = get_ip()
+    fsdp_master_port = get_open_port()
+
+    # Launch 4 FSDP training workers.
+    # Ray allocates 1 GPU per worker; AsyncLLMEngine's internal DP
+    # placement groups will land on the remaining 4 GPUs.
+    fsdp_workers = [
+        FSDPTrainWorker.remote(
+            local_model_path,
+            rank,
+            FSDP_WORLD_SIZE,
+            fsdp_master_addr,
+            fsdp_master_port,
+        )
+        for rank in range(FSDP_WORLD_SIZE)
+    ]
+    ray.get([w.get_rank.remote() for w in fsdp_workers])
+    print(f"[init] {FSDP_WORLD_SIZE} FSDP training workers ready.")
+
+    # Launch vLLM with expert parallelism + data parallelism.
+    # AsyncLLMEngine with data_parallel_backend="ray" creates its own
+    # placement groups internally — no manual placement group needed.
+    print("[engine] Creating AsyncLLMEngine...")
+    engine = create_async_engine(
+        model=local_model_path,
+        enforce_eager=True,
+        tensor_parallel_size=INFERENCE_TP_SIZE,
+        data_parallel_size=INFERENCE_DP_SIZE,
+        enable_expert_parallel=True,
+        distributed_executor_backend="ray",
+        data_parallel_backend="ray",
+        weight_transfer_config=WeightTransferConfig(backend="nccl"),
+        load_format="dummy",
+        gpu_memory_utilization=0.7,
+    )
+    print("[engine] AsyncLLMEngine created.")
+
+    prompts = [
+        "Hello, my name is",
+        "The president of the United States is",
+        "The capital of France is",
+        "The future of AI is",
+    ]
+    sampling_params = SamplingParams(temperature=0)
+
+    # Generate with dummy weights — expect gibberish.
+    print("[generate] Starting generation with dummy weights...")
+    outputs = await generate_batch(engine, prompts, sampling_params)
+    print("[generate] Generation complete.")
+
+    print("-" * 60)
+    print("BEFORE weight sync (dummy weights):")
+    print("-" * 60)
+    for output in outputs:
+        print(f"Prompt: {output.prompt!r}")
+        print(f"Generated: {output.outputs[0].text!r}")
+        print("-" * 60)
+
+    # --- Weight-transfer setup ---
+    print("[transfer] Setting up weight-transfer endpoint...")
+    transfer_addr, transfer_port = ray.get(
+        fsdp_workers[0].setup_transfer_endpoint.remote()
+    )
+    print(f"[transfer] Endpoint ready at {transfer_addr}:{transfer_port}")
+
+    transfer_world_size = INFERENCE_TP_SIZE * INFERENCE_DP_SIZE + 1
+    print(
+        f"[transfer] World size: {transfer_world_size} "
+        f"(1 trainer + {INFERENCE_TP_SIZE * INFERENCE_DP_SIZE} vLLM workers)"
+    )
+
+    print("[transfer] Initializing NCCL groups...")
+    train_handle = fsdp_workers[0].init_weight_transfer_group.remote(
+        transfer_world_size
+    )
+    await engine.init_weight_transfer_engine(
+        WeightTransferInitRequest(
+            init_info=asdict(
+                NCCLWeightTransferInitInfo(
+                    master_address=transfer_addr,
+                    master_port=transfer_port,
+                    rank_offset=1,
+                    world_size=transfer_world_size,
+                )
+            )
+        )
+    )
+    ray.get(train_handle)
+    print("[transfer] NCCL groups initialized.")
+
+    # --- Pause, transfer weights, resume ---
+    print("[sync] Pausing generation...")
+    await engine.pause_generation(mode="abort")
+    print("[sync] Generation paused.")
+
+    names, dtype_names, shapes = ray.get(fsdp_workers[0].get_weight_metadata.remote())
+    print(f"[sync] Got metadata for {len(names)} parameters.")
+
+    print("[sync] Broadcasting weights from FSDP → vLLM...")
+    broadcast_handles = [
+        w.gather_and_broadcast_weights.remote(packed=True) for w in fsdp_workers
+    ]
+    await engine.update_weights(
+        WeightTransferUpdateRequest(
+            update_info=asdict(
+                NCCLWeightTransferUpdateInfo(
+                    names=names,
+                    dtype_names=dtype_names,
+                    shapes=shapes,
+                    packed=True,
+                )
+            )
+        )
+    )
+    ray.get(broadcast_handles)
+    print("[sync] Weight broadcast complete.")
+
+    print("[sync] Resuming generation...")
+    await engine.resume_generation()
+    print("[sync] Generation resumed.")
+
+    # Generate with synced weights — expect sensible output.
+    print("[generate] Starting generation with synced weights...")
+    outputs_updated = await generate_batch(engine, prompts, sampling_params)
+    print("[generate] Generation complete.")
+
+    print("-" * 60)
+    print("AFTER weight sync (real weights):")
+    print("-" * 60)
+    for output in outputs_updated:
+        print(f"Prompt: {output.prompt!r}")
+        print(f"Generated: {output.outputs[0].text!r}")
+        print("-" * 60)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
@@ -5,6 +5,7 @@ import numpy as np
 import pytest
 import torch

+from vllm.distributed.eplb.eplb_state import compute_logical_maps
 from vllm.distributed.eplb.policy.default import DefaultEplbPolicy


@@ -24,9 +25,10 @@ def test_basic_rebalance():
    num_nodes = 2
    num_gpus = 8

-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
        weight, num_replicas, num_groups, num_nodes, num_gpus
    )
+    log2phy, logcnt = compute_logical_maps(phy2log, weight.shape[-1])

    # Verify output shapes
    assert phy2log.shape == (
@@ -78,9 +80,10 @@ def test_single_gpu_case():
    num_nodes = 1
    num_gpus = 1

-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
        weight, num_replicas, num_groups, num_nodes, num_gpus
    )
+    log2phy, logcnt = compute_logical_maps(phy2log, weight.shape[-1])

    # Verify shapes
    assert phy2log.shape == (1, 4)
@@ -100,9 +103,10 @@ def test_equal_weights():
    num_nodes = 2
    num_gpus = 4

-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
        weight, num_replicas, num_groups, num_nodes, num_gpus
    )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])

    # Verify shapes
    assert phy2log.shape == (1, 8)
@@ -123,9 +127,10 @@ def test_extreme_weight_imbalance():
    num_nodes = 2
    num_gpus = 4

-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
        weight, num_replicas, num_groups, num_nodes, num_gpus
    )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])

    # Verify shapes
    assert phy2log.shape == (1, 12)
@@ -151,9 +156,10 @@ def test_multiple_layers():
    num_nodes = 2
    num_gpus = 4

-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
        weight, num_replicas, num_groups, num_nodes, num_gpus
    )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])

    # Verify shapes
    assert phy2log.shape == (3, 8)
@@ -176,7 +182,8 @@ def test_parameter_validation():
    # Test non-divisible case - this should handle normally without throwing
    # errors because the function will fall back to global load balancing
    # strategy
-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(weight, 8, 3, 2, 4)
+    phy2log = DefaultEplbPolicy.rebalance_experts(weight, 8, 3, 2, 4)
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])
    assert phy2log.shape == (1, 8)
    assert logcnt.shape == (1, 4)

@@ -198,9 +205,10 @@ def test_small_scale_hierarchical():
    num_nodes = 2  # 2 nodes
    num_gpus = 4  # 4 GPUs

-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
        weight, num_replicas, num_groups, num_nodes, num_gpus
    )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])

    # Verify basic constraints
    assert phy2log.shape == (1, 12)
@@ -225,9 +233,10 @@ def test_global_load_balance_fallback():
    num_nodes = 2
    num_gpus = 4

-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
        weight, num_replicas, num_groups, num_nodes, num_gpus
    )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])

    # Should work normally, just using global load balancing strategy
    assert phy2log.shape == (1, 8)
@@ -247,9 +256,10 @@ def test_device_compatibility(device):
    num_nodes = 1
    num_gpus = 2

-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
        weight, num_replicas, num_groups, num_nodes, num_gpus
    )
+    _, logcnt = compute_logical_maps(phy2log, weight.shape[-1])

    # Function will convert to CPU internally, but should handle different
    # device inputs normally
@@ -264,9 +274,8 @@ def test_additional_cases():
    weight1 = torch.tensor(
        [[50, 100, 75, 120, 90, 60, 80, 110, 40, 70, 95, 85, 65, 55, 45, 35]]
    )
-    phy2log1, log2phy1, logcnt1 = DefaultEplbPolicy.rebalance_experts(
-        weight1, 24, 8, 4, 8
-    )
+    phy2log1 = DefaultEplbPolicy.rebalance_experts(weight1, 24, 8, 4, 8)
+    _, logcnt1 = compute_logical_maps(phy2log1, weight1.shape[-1])

    assert phy2log1.shape == (1, 24)
    assert logcnt1.shape == (1, 16)
@@ -279,9 +288,8 @@ def test_additional_cases():
            [12, 25, 50, 100, 150, 200],  # Increasing weights
        ]
    )
-    phy2log2, log2phy2, logcnt2 = DefaultEplbPolicy.rebalance_experts(
-        weight2, 10, 3, 1, 2
-    )
+    phy2log2 = DefaultEplbPolicy.rebalance_experts(weight2, 10, 3, 1, 2)
+    _, logcnt2 = compute_logical_maps(phy2log2, weight2.shape[-1])

    assert phy2log2.shape == (2, 10)
    assert logcnt2.shape == (2, 6)
@@ -292,6 +300,42 @@ def test_additional_cases():
        assert logcnt2[layer, max_weight_idx] >= 2


+def test_compute_logical_maps_with_negative_indices():
+    """
+    Test that compute_logical_maps correctly handles physical slots containing
+    -1 (unused slots).
+    """
+    # 2 layers, 6 physical slots, 4 logical experts.
+    # Slots 2 and 5 are unused (-1).
+    phy2log = torch.tensor(
+        [
+            [0, 1, -1, 2, 3, -1],
+            [3, -1, 2, 1, 0, -1],
+        ]
+    )
+    num_layers = 2
+    num_logical_experts = 4
+
+    log2phy, logcnt = compute_logical_maps(phy2log, num_logical_experts)
+
+    assert logcnt.shape == (num_layers, num_logical_experts)
+    assert log2phy.shape == (num_layers, num_logical_experts, 1)
+
+    expected_logcnt = torch.ones(num_layers, num_logical_experts, dtype=phy2log.dtype)
+    assert torch.all(logcnt == expected_logcnt), (
+        f"Expected that all replica counts == 1, got {logcnt}"
+    )
+
+    assert torch.all(log2phy >= 0), (
+        "log2phy should only contain valid physical indices, not -1"
+    )
+
+    assert log2phy[0, 0, 0] == 0
+    assert log2phy[0, 1, 0] == 1
+    assert log2phy[0, 2, 0] == 3
+    assert log2phy[0, 3, 0] == 4
+
+
 if __name__ == "__main__":
    weight = torch.tensor(
        [
@@ -305,7 +349,7 @@ if __name__ == "__main__":
    num_nodes = 2
    num_gpus = 8

-    phy2log, log2phy, logcnt = DefaultEplbPolicy.rebalance_experts(
+    phy2log = DefaultEplbPolicy.rebalance_experts(
        weight, num_replicas, num_groups, num_nodes, num_gpus
    )
    print(phy2log)
@@ -434,9 +478,10 @@ def test_preserve_intragpu_slots(
    """Experts that stay on a GPU keep their old slots; incoming not lost."""
    phy_replicas_idx = _make_phy_replicas_idx_from_phy2log(new_phy2log)

-    post_phy2log, post_phy_replicas_idx = DefaultEplbPolicy.preserve_intragpu_slots(
-        new_phy2log, phy_replicas_idx, num_ranks, old_phy2log
+    post_phy2log = DefaultEplbPolicy.preserve_intragpu_slots(
+        new_phy2log, num_ranks, old_phy2log
    )
+    post_phy_replicas_idx = _make_phy_replicas_idx_from_phy2log(post_phy2log)

    # Shapes preserved
    assert post_phy2log.shape == new_phy2log.shape
@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project


+import itertools
+
 import pytest
 import torch

@@ -18,17 +20,17 @@ from vllm.platforms import current_platform

 DTYPES = [torch.bfloat16, torch.float]
 QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
-# Trimmed to cover: small, misaligned, large-aligned, large-misaligned
+VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
+# Avoid combinatorial explosion with full Cartesian product
 NUM_TOKENS_HIDDEN_SIZES = [
-    (1, 128),
-    (1, 1025),       # odd/misaligned vectorization
-    (2048, 1024),     # medium aligned
-    (4096, 5137),     # large misaligned
+    *[(1, i) for i in [1, 64, 128, *VEC_HIDDEN_SIZES, 5120, 5137]],
+    *[(2048, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5137]],
+    *[(4096, i) for i in [1, 64, 5137]],
 ]

 ADD_RESIDUAL = [False, True]
 SCALE_UBS = [True, False]
-GROUP_SIZES = [None, [1, 128]]
+GROUP_SIZES = [None, [1, 64], [1, 128]]
 TMA_ALIGNMENTS = [0, 4]
 SEEDS = [0]
 CUDA_DEVICES = [
@@ -158,7 +160,7 @@ def ops_impl(
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
@pytest.mark.parametrize(
    "group_size, tma_alignment",
-    [(None, 0), ([1, 128], 0), ([1, 128], 4)],
+    [(None, 0), *itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
 )
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@@ -10,8 +10,8 @@ from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.utils.torch_utils import set_random_seed

 DTYPES = [torch.half, torch.bfloat16, torch.float]
-NUM_TOKENS = [7, 4096]  # Small + large
-HIDDEN_SIZES = [8, 769, 8192]  # Small, odd/misaligned, large
+NUM_TOKENS = [7, 83, 4096]  # Arbitrary values for testing
+HIDDEN_SIZES = [8, 768, 769, 5120, 5125, 8192]  # Arbitrary values for testing
 ADD_RESIDUAL = [False, True]
 SEEDS = [0]
 CUDA_DEVICES = [
@@ -77,7 +77,7 @@ def test_rms_norm(
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
@pytest.mark.parametrize("dtype", DTYPES)
-@pytest.mark.parametrize("quant_scale", [0.01, 10.0])
+@pytest.mark.parametrize("quant_scale", [0.01, 1.0, 10.0])
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("strided_input", [False, True])
@@ -1,165 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for zero-overhead NaN/Inf detection in RMSNorm kernels."""
-
-import pytest
-import torch
-
-from vllm.model_executor.layers.nan_detector import NaNDetector
-
-
-@pytest.fixture(autouse=True)
-def reset_nan_detector():
-    """Reset the singleton between tests."""
-    NaNDetector.reset()
-    yield
-    NaNDetector.reset()
-
-
-@pytest.fixture
-def device():
-    return "cuda:0"
-
-
-@pytest.mark.parametrize("hidden_size", [64, 128, 256])
-@pytest.mark.parametrize("num_tokens", [1, 4, 16])
-@torch.inference_mode()
-def test_nan_detection_rms_norm(default_vllm_config, device, hidden_size, num_tokens):
-    """NaN in input should be detected at the correct token position."""
-    from vllm import _custom_ops as ops
-
-    num_layers = 3
-    max_num_tokens = 32
-
-    nan_flags = torch.zeros(num_layers, max_num_tokens, dtype=torch.int8, device=device)
-    weight = torch.ones(hidden_size, dtype=torch.float16, device=device)
-
-    # Clean input — no flags should be set.
-    x = torch.randn(num_tokens, hidden_size, dtype=torch.float16, device=device)
-    out = torch.empty_like(x)
-    ops.rms_norm(out, x, weight, 1e-6, nan_flags, 0, max_num_tokens)
-    assert nan_flags.sum().item() == 0, "False positive on clean input"
-
-    # Inject NaN at token 1, layer 0.
-    nan_flags.zero_()
-    x_nan = x.clone()
-    if num_tokens > 1:
-        x_nan[1, 0] = float("nan")
-        ops.rms_norm(out, x_nan, weight, 1e-6, nan_flags, 0, max_num_tokens)
-        assert nan_flags[0, 1].item() == 1, "NaN not detected at token 1"
-        assert nan_flags[0, 0].item() == 0, "False positive at token 0"
-    else:
-        x_nan[0, 0] = float("nan")
-        ops.rms_norm(out, x_nan, weight, 1e-6, nan_flags, 0, max_num_tokens)
-        assert nan_flags[0, 0].item() == 1, "NaN not detected at token 0"
-
-    # Inject NaN at a different layer index.
-    nan_flags.zero_()
-    ops.rms_norm(out, x_nan, weight, 1e-6, nan_flags, 2, max_num_tokens)
-    assert nan_flags[0].sum().item() == 0, "Wrong layer got the flag"
-    assert nan_flags[2].any().item(), "NaN not detected at layer 2"
-
-
-@pytest.mark.parametrize("hidden_size", [64, 256])
-@torch.inference_mode()
-def test_inf_detection_rms_norm(default_vllm_config, device, hidden_size):
-    """Inf in input should be detected."""
-    from vllm import _custom_ops as ops
-
-    num_tokens = 4
-    max_num_tokens = 8
-    nan_flags = torch.zeros(1, max_num_tokens, dtype=torch.int8, device=device)
-    weight = torch.ones(hidden_size, dtype=torch.float16, device=device)
-
-    x = torch.randn(num_tokens, hidden_size, dtype=torch.float16, device=device)
-    x[2, 0] = float("inf")
-    out = torch.empty_like(x)
-    ops.rms_norm(out, x, weight, 1e-6, nan_flags, 0, max_num_tokens)
-    assert nan_flags[0, 2].item() == 1, "Inf not detected at token 2"
-
-
-@pytest.mark.parametrize("hidden_size", [64, 256])
-@torch.inference_mode()
-def test_nan_detection_fused_add_rms_norm(default_vllm_config, device, hidden_size):
-    """NaN detection works with the fused add+norm path."""
-    from vllm import _custom_ops as ops
-
-    num_tokens = 4
-    max_num_tokens = 8
-    nan_flags = torch.zeros(1, max_num_tokens, dtype=torch.int8, device=device)
-    weight = torch.ones(hidden_size, dtype=torch.float16, device=device)
-
-    x = torch.randn(num_tokens, hidden_size, dtype=torch.float16, device=device)
-    residual = torch.randn_like(x)
-
-    # Clean — no flags.
-    ops.fused_add_rms_norm(
-        x.clone(), residual.clone(), weight, 1e-6, nan_flags, 0, max_num_tokens
-    )
-    assert nan_flags.sum().item() == 0
-
-    # Inject NaN in the input (not residual).
-    nan_flags.zero_()
-    x_nan = x.clone()
-    x_nan[3, 0] = float("nan")
-    ops.fused_add_rms_norm(
-        x_nan, residual.clone(), weight, 1e-6, nan_flags, 0, max_num_tokens
-    )
-    assert nan_flags[0, 3].item() == 1, "NaN not detected at token 3"
-
-    # Inject NaN in the residual.
-    nan_flags.zero_()
-    res_nan = residual.clone()
-    res_nan[0, 0] = float("nan")
-    ops.fused_add_rms_norm(
-        x.clone(), res_nan, weight, 1e-6, nan_flags, 0, max_num_tokens
-    )
-    assert nan_flags[0, 0].item() == 1, "NaN in residual not detected"
-
-
-@torch.inference_mode()
-def test_no_detection_when_disabled(default_vllm_config, device):
-    """When nan_flags is None, no detection occurs (null pointer path)."""
-    from vllm import _custom_ops as ops
-
-    hidden_size = 64
-    num_tokens = 4
-    weight = torch.ones(hidden_size, dtype=torch.float16, device=device)
-    x = torch.randn(num_tokens, hidden_size, dtype=torch.float16, device=device)
-    x[0, 0] = float("nan")
-    out = torch.empty_like(x)
-
-    # Should not crash — nan_flags=None means no detection.
-    ops.rms_norm(out, x, weight, 1e-6)
-
-
-@torch.inference_mode()
-def test_nan_detector_class(default_vllm_config, device):
-    """Test the NaNDetector singleton lifecycle."""
-    detector = NaNDetector.get()
-
-    # Register layers.
-    idx0 = detector.register("layer_0")
-    idx1 = detector.register("layer_1")
-    assert idx0 == 0
-    assert idx1 == 1
-
-    # Finalize.
-    max_tokens = 8
-    detector.finalize(torch.device(device), max_tokens)
-    assert detector.nan_flags is not None
-    assert detector.nan_flags.shape == (2, max_tokens)
-    assert detector.max_num_tokens == max_tokens
-
-    # Clear + check with no NaN — should log nothing.
-    detector.clear()
-    detector.check(4)  # 4 real tokens
-
-    # Manually set a flag and check.
-    detector.nan_flags[0, 2] = 1
-    detector.check(4)  # Should log ERROR for layer_0, token 2
-
-    # Set a flag in padding region.
-    detector.clear()
-    detector.nan_flags[1, 6] = 1
-    detector.check(4)  # Should log WARNING for layer_1 (padding)
@@ -1,134 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test for NVFP4 NaN propagation within a SINGLE TOKEN when NaN appears in some
-feature dimensions but not others.
-
-This is the REAL bug: if a single token has NaN in some dimensions (e.g., from
-a buggy attention output), the block scale for that block becomes NaN, which
-then contaminates the ENTIRE output for that token.
-"""
-import pytest
-import torch
-
-from vllm import _custom_ops as ops
-from vllm.platforms import current_platform
-from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer
-
-if not current_platform.has_device_capability(100):
-    pytest.skip(
-        reason="NVFP4 requires compute capability 100 or above (Blackwell+).",
-        allow_module_level=True,
-    )
-
-if not has_flashinfer():
-    pytest.skip(
-        reason="FlashInfer is required for NVFP4 GEMM tests.",
-        allow_module_level=True,
-    )
-
-FLOAT4_E2M1_MAX = 6.0
-FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
-
-
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("use_fix", [True, False])
-@torch.inference_mode()
-def test_nvfp4_nan_within_token_contamination(dtype: torch.dtype, use_fix: bool) -> None:
-    """
-    Test that NaN in a few dimensions of a token contaminates the entire token output.
-
-    Setup:
-    - Single token with mostly clean values
-    - NaN injected into ONE BLOCK of the token (e.g., dimensions 16-31)
-    - This makes that block's scale = NaN
-    - The entire token output becomes NaN (not just the output dimensions
-      corresponding to that block)
-    """
-    device = "cuda:0"
-    torch.set_default_device(device)
-    torch.manual_seed(42)
-
-    # Single token with hidden_size=64 (4 blocks of 16)
-    x = torch.randn(1, 64, dtype=dtype, device=device)
-
-    # Inject NaN into the SECOND BLOCK (dims 16-31) of this token
-    x[0, 16:32] = float('nan')
-
-    print(f"\nInput token:")
-    print(f"  Block 0 (dims 0-15): clean, sample={x[0, 0:4]}")
-    print(f"  Block 1 (dims 16-31): NaN, sample={x[0, 16:20]}")
-    print(f"  Block 2 (dims 32-47): clean, sample={x[0, 32:36]}")
-    print(f"  Block 3 (dims 48-63): clean, sample={x[0, 48:52]}")
-
-    # Apply fix if requested
-    if use_fix:
-        x = torch.where(torch.isnan(x), torch.zeros_like(x), x)
-        print(f"\n[FIX APPLIED] NaNs masked to zero")
-
-    # Compute global scale
-    input_amax = torch.abs(x[torch.isfinite(x)]).max().to(torch.float32)
-    input_global_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / input_amax).to(torch.float32)
-    input_global_scale_inv = 1.0 / input_global_scale
-
-    # Quantize
-    x_fp4, x_blockscale = ops.scaled_fp4_quant(
-        x, input_global_scale_inv, is_sf_swizzled_layout=False,
-        backend="flashinfer-cutlass")
-
-    print(f"\nBlock scales after quantization:")
-    for i in range(4):
-        scale_val = x_blockscale.view(torch.float8_e4m3fn)[0, i].to(torch.float32)
-        print(f"  Block {i}: {scale_val}")
-
-    # Create weights
-    output_size = 128
-    weight = torch.randn(output_size, 64, dtype=dtype, device=device)
-    weight_amax = torch.abs(weight).max().to(torch.float32)
-    weight_global_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / weight_amax).to(torch.float32)
-    weight_fp4, weight_blockscale = ops.scaled_fp4_quant(
-        weight, weight_global_scale, is_sf_swizzled_layout=False)
-
-    alpha = (input_global_scale * weight_global_scale).to(torch.float32)
-
-    # Run GEMM
-    output = flashinfer_scaled_fp4_mm(
-        x_fp4, weight_fp4, x_blockscale, weight_blockscale, alpha, dtype,
-        backend="cutlass")
-
-    print(f"\nOutput shape: {output.shape}")
-    print(f"Output sample (first 8): {output[0, :8]}")
-
-    has_nan = torch.isnan(output).any()
-    print(f"Has NaN in output: {has_nan}")
-
-    if has_nan:
-        nan_percentage = 100.0 * torch.isnan(output).sum().item() / output.numel()
-        print(f"NaN percentage: {nan_percentage:.1f}%")
-
-        if use_fix:
-            pytest.fail(
-                f"NaN contamination detected even with fix applied!\n"
-                f"  {nan_percentage:.1f}% of output is NaN\n"
-                f"  The fix should have prevented this."
-            )
-        else:
-            pytest.fail(
-                f"NaN contamination detected (expected on buggy path)!\n"
-                f"  A single NaN block in the input caused {nan_percentage:.1f}% of output to be NaN\n"
-                f"  This demonstrates the bug: NaN in one block contaminates the entire token output."
-            )
-    else:
-        print("✓ No NaN contamination detected")
-
-
-if __name__ == "__main__":
-    print("="*60)
-    print("Testing BUGGY PATH (no NaN masking)")
-    print("="*60)
-    test_nvfp4_nan_within_token_contamination(dtype=torch.bfloat16, use_fix=False)
-
-    print("\n" + "="*60)
-    print("Testing FIXED PATH (with NaN masking)")
-    print("="*60)
-    test_nvfp4_nan_within_token_contamination(dtype=torch.bfloat16, use_fix=True)
@@ -1,125 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Integration test for NVFP4 NaN handling through the full apply_nvfp4_linear path.
-
-This verifies that the production code fix in apply_nvfp4_linear() properly
-masks NaNs before quantization.
-"""
-import pytest
-import torch
-import torch.nn as nn
-
-from vllm import _custom_ops as ops
-from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
-    apply_nvfp4_linear,
-    convert_to_nvfp4_linear_kernel_format,
-    select_nvfp4_linear_backend,
-)
-from vllm.platforms import current_platform
-from vllm.utils.flashinfer import has_flashinfer
-
-if not current_platform.has_device_capability(100):
-    pytest.skip(
-        reason="NVFP4 requires compute capability 100 or above (Blackwell+).",
-        allow_module_level=True,
-    )
-
-if not has_flashinfer():
-    pytest.skip(
-        reason="FlashInfer is required for NVFP4 tests.",
-        allow_module_level=True,
-    )
-
-FLOAT4_E2M1_MAX = 6.0
-FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
-
-
-def create_nvfp4_layer(input_size: int, output_size: int, dtype: torch.dtype,
-                       device: str) -> tuple[nn.Module, any]:
-    """Create a mock NVFP4 linear layer for testing."""
-    layer = nn.Module()
-    layer.input_size_per_partition = input_size
-    layer.output_size_per_partition = output_size
-
-    # Create and quantize random weights
-    weight_bf16 = torch.randn(output_size, input_size, dtype=dtype, device=device)
-    weight_amax = torch.abs(weight_bf16).max().to(torch.float32)
-    weight_global_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / weight_amax).to(torch.float32)
-
-    # Quantize weights to FP4
-    weight_fp4, weight_blockscale = ops.scaled_fp4_quant(
-        weight_bf16, weight_global_scale, is_sf_swizzled_layout=True,
-        backend="flashinfer-cutlass")
-
-    layer.weight = nn.Parameter(weight_fp4, requires_grad=False)
-    layer.weight_scale = nn.Parameter(weight_blockscale, requires_grad=False)
-
-    # Global scales
-    layer.weight_global_scale = nn.Parameter(weight_global_scale, requires_grad=False)
-
-    # Input scale (will be computed per-batch, this is just placeholder)
-    input_global_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
-    layer.input_global_scale_inv = nn.Parameter(1.0 / input_global_scale, requires_grad=False)
-    layer.alpha = nn.Parameter(input_global_scale * weight_global_scale, requires_grad=False)
-
-    # Convert to kernel format
-    backend = select_nvfp4_linear_backend()
-    convert_to_nvfp4_linear_kernel_format(backend, layer)
-
-    return layer, backend
-
-
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@torch.inference_mode()
-def test_nvfp4_linear_with_nan_input(dtype: torch.dtype) -> None:
-    """
-    Test that apply_nvfp4_linear handles NaN inputs correctly.
-
-    This is an end-to-end integration test using the production code path.
-    """
-    device = "cuda:0"
-    torch.set_default_device(device)
-    torch.manual_seed(42)
-
-    input_size = 64
-    output_size = 128
-    batch_size = 4
-
-    # Create layer
-    layer, backend = create_nvfp4_layer(input_size, output_size, dtype, device)
-
-    # Create input with NaN in some positions
-    x = torch.randn(batch_size, input_size, dtype=dtype, device=device)
-
-    # Inject NaN into token 2, block 1 (dimensions 16-31)
-    x[2, 16:32] = float('nan')
-
-    print(f"\nInput shape: {x.shape}")
-    print(f"Token 2, block 1 has NaN: {torch.isnan(x[2, 16:32]).all()}")
-    print(f"Other tokens clean: {not torch.isnan(x[[0,1,3]]).any()}")
-
-    # Apply NVFP4 linear (production code path with fix)
-    output = apply_nvfp4_linear(backend=backend, layer=layer, x=x, bias=None)
-
-    print(f"\nOutput shape: {output.shape}")
-    print(f"Output token 0 (clean input): {output[0, :8]}")
-    print(f"Output token 2 (had NaN input): {output[2, :8]}")
-
-    # Check results
-    has_nan = torch.isnan(output).any()
-    print(f"\nHas NaN in output: {has_nan}")
-
-    if has_nan:
-        nan_percentage = 100.0 * torch.isnan(output).sum().item() / output.numel()
-        pytest.fail(
-            f"NaN detected in output!\n"
-            f"  {nan_percentage:.1f}% of output is NaN\n"
-            f"  The fix in apply_nvfp4_linear should have masked NaNs before quantization."
-        )
-
-    print("✓ No NaN in output - fix is working correctly!")
-
-
-if __name__ == "__main__":
-    test_nvfp4_linear_with_nan_input(dtype=torch.bfloat16)
@@ -1,375 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""
-Test for NVFP4 GEMM NaN propagation from padding positions.
-
-This test validates that NaNs in padding positions (from attention softmax 0/0)
-do not leak into real token positions during FP4 quantization and GEMM.
-"""
-import pytest
-import torch
-
-from vllm import _custom_ops as ops
-from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
-    NvFp4LinearBackend,
-    pad_nvfp4_activation_for_cutlass,
-    pad_nvfp4_weight_for_cutlass,
-    slice_nvfp4_output,
-    swizzle_blockscale,
-)
-from vllm.platforms import current_platform
-from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer
-from vllm.utils.torch_utils import set_random_seed
-
-if not current_platform.has_device_capability(100):
-    pytest.skip(
-        reason="NVFP4 requires compute capability 100 or above (Blackwell+).",
-        allow_module_level=True,
-    )
-
-if not has_flashinfer():
-    pytest.skip(
-        reason="FlashInfer is required for NVFP4 GEMM tests.",
-        allow_module_level=True,
-    )
-
-FLOAT4_E2M1_MAX = 6.0
-FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
-
-
-def create_nvfp4_weight(output_size: int, input_size: int, dtype: torch.dtype,
-                        device: str) -> tuple[torch.Tensor, torch.Tensor, float, int]:
-    """Create random FP4 weights and scales for testing."""
-    # Create random bf16 weights
-    weight_bf16 = torch.randn(output_size, input_size, dtype=dtype, device=device)
-
-    # Compute global scale
-    weight_amax = torch.abs(weight_bf16).max().to(torch.float32)
-    weight_global_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / weight_amax).to(
-        torch.float32)
-
-    # Quantize to FP4
-    weight_fp4, weight_blockscale = ops.scaled_fp4_quant(
-        weight_bf16, weight_global_scale, is_sf_swizzled_layout=True)
-
-    # Swizzle block scales for CUTLASS kernel
-    weight_scale_swizzled = swizzle_blockscale(
-        weight_blockscale.view(torch.float8_e4m3fn))
-
-    # Pad weight for CUTLASS alignment
-    weight_fp4_padded, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
-        weight_fp4)
-
-    return weight_fp4_padded, weight_scale_swizzled, weight_global_scale, weights_padding_cols
-
-
-@pytest.mark.parametrize("num_tokens", [32])
-@pytest.mark.parametrize("num_padding", [8])
-@pytest.mark.parametrize("hidden_size", [1024])
-@pytest.mark.parametrize("output_size", [1024])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("nan_placement", ["end"])
-@pytest.mark.parametrize("use_buggy_path", [True, False])
-@torch.inference_mode()
-def test_nvfp4_gemm_nan_isolation(
-    num_tokens: int,
-    num_padding: int,
-    hidden_size: int,
-    output_size: int,
-    dtype: torch.dtype,
-    nan_placement: str,
-    use_buggy_path: bool,
-) -> None:
-    """
-    Test that NaNs in padding positions don't leak into real token positions.
-
-    Simulates the scenario where attention softmax produces NaN at padding
-    positions (0/0), which then flows through o_proj's NVFP4 GEMM.
-
-    Args:
-        num_tokens: Number of real (non-padding) tokens
-        num_padding: Number of padding tokens with NaN
-        hidden_size: Input dimension (K)
-        output_size: Output dimension (N)
-        dtype: Input data type
-        nan_placement: Where to place NaN tokens ("end", "middle", "scattered")
-        use_buggy_path: If True, don't mask NaNs before quantization (buggy).
-                        If False, mask NaNs before quantization (fixed).
-    """
-    set_random_seed(42)
-    device = "cuda:0"
-    torch.set_default_device(device)
-
-    total_tokens = num_tokens + num_padding
-
-    # Create input with NaNs at padding positions
-    x = torch.randn(total_tokens, hidden_size, dtype=dtype, device=device)
-
-    # Create a mask: 1 for real tokens, 0 for padding
-    mask = torch.ones(total_tokens, dtype=torch.bool, device=device)
-
-    # Inject NaNs at padding positions based on placement strategy
-    if nan_placement == "end":
-        # NaNs at the end (most common case)
-        x[num_tokens:, :] = float('nan')
-        mask[num_tokens:] = False
-    elif nan_placement == "middle":
-        # NaNs in the middle
-        mid_start = num_tokens // 2
-        x[mid_start:mid_start + num_padding, :] = float('nan')
-        mask[mid_start:mid_start + num_padding] = False
-    elif nan_placement == "scattered":
-        # Scattered NaN positions
-        nan_indices = torch.randperm(total_tokens)[:num_padding]
-        x[nan_indices, :] = float('nan')
-        mask[nan_indices] = False
-
-    # Verify NaNs are present at padding positions
-    assert torch.isnan(x[~mask]).all(), "NaN injection failed"
-    assert not torch.isnan(x[mask]).any(), "Real tokens should not have NaN"
-
-    # Create FP4 weights
-    weight_fp4, weight_scale, weight_global_scale, weights_padding_cols = \
-        create_nvfp4_weight(output_size, hidden_size, dtype, device)
-
-    # Compute input global scale
-    # Always use clean tokens for global scale (even in buggy path)
-    # because NaN global scale would make everything NaN
-    input_amax = torch.abs(x[mask]).max().to(torch.float32)
-    input_global_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / input_amax).to(
-        torch.float32)
-    input_global_scale_inv = 1.0 / input_global_scale
-    alpha = (input_global_scale * weight_global_scale).to(torch.float32)
-
-    # **KEY DIFFERENCE**: Buggy path vs fixed path
-    if use_buggy_path:
-        # BUGGY: Pass input with NaNs directly to quantization
-        # This allows NaNs to contaminate block scales
-        x_to_quantize = x
-    else:
-        # FIXED: Mask NaNs before quantization
-        # This prevents NaNs from contaminating block scales
-        x_to_quantize = torch.where(torch.isnan(x), torch.zeros_like(x), x)
-
-    # Quantize input to FP4 (this is where NaN propagation can happen)
-    x_fp4, x_blockscale = ops.scaled_fp4_quant(
-        x_to_quantize, input_global_scale_inv, is_sf_swizzled_layout=True,
-        backend="flashinfer-cutlass")
-
-    # Pad activations to match weight K-dimension padding
-    x_fp4_padded = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
-
-    # Run the FP4 GEMM (FlashInfer CUTLASS backend)
-    output = flashinfer_scaled_fp4_mm(
-        x_fp4_padded,
-        weight_fp4,
-        x_blockscale,
-        weight_scale,
-        alpha,
-        dtype,
-        backend="cutlass",
-    )
-
-    # Slice output to remove N-dimension padding
-    output = slice_nvfp4_output(output, output_size)
-
-    # Check for NaN propagation
-    real_output = output[mask]  # Real token outputs
-    padding_output = output[~mask]  # Padding token outputs
-
-    has_nan_in_real = torch.isnan(real_output).any()
-    has_nan_in_padding = torch.isnan(padding_output).any()
-
-    # Collect statistics for debugging
-    if has_nan_in_real:
-        num_nan_elements = torch.isnan(real_output).sum().item()
-        total_real_elements = real_output.numel()
-        nan_percentage = 100.0 * num_nan_elements / total_real_elements
-
-        if use_buggy_path:
-            # Expected to fail on buggy path - this confirms the bug exists
-            pytest.fail(
-                f"NaN LEAK DETECTED (buggy path - expected to fail)!\n"
-                f"  Configuration: {num_tokens} real + {num_padding} padding tokens\n"
-                f"  NaN placement: {nan_placement}\n"
-                f"  Input shape: {x.shape}, Output shape: {output.shape}\n"
-                f"  NaN in real output: {num_nan_elements}/{total_real_elements} "
-                f"({nan_percentage:.2f}%)\n"
-                f"  NaN in padding output: {has_nan_in_padding}\n"
-                f"This confirms the hypothesis that NaNs leak from padding to real tokens."
-            )
-        else:
-            # Should NOT fail on fixed path
-            pytest.fail(
-                f"NaN LEAK DETECTED (fixed path - should not happen)!\n"
-                f"  The fix (NaN masking) did not work as expected.\n"
-                f"  Configuration: {num_tokens} real + {num_padding} padding tokens\n"
-                f"  NaN placement: {nan_placement}\n"
-                f"  NaN in real output: {num_nan_elements}/{total_real_elements} "
-                f"({nan_percentage:.2f}%)"
-            )
-
-    # If we reach here, NaNs are properly isolated
-    path_type = "buggy" if use_buggy_path else "fixed"
-    print(f"✓ NaN isolation verified ({path_type} path): {nan_placement} placement, "
-          f"{num_tokens} real + {num_padding} padding tokens")
-
-
-@pytest.mark.parametrize("num_tokens", [32])
-@pytest.mark.parametrize("num_padding", [8])
-@pytest.mark.parametrize("hidden_size", [1024])
-@pytest.mark.parametrize("output_size", [1024])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@torch.inference_mode()
-def test_nvfp4_gemm_nan_masking_fix(
-    num_tokens: int,
-    num_padding: int,
-    hidden_size: int,
-    output_size: int,
-    dtype: torch.dtype,
-) -> None:
-    """
-    Test a potential fix: masking NaNs before FP4 quantization.
-
-    This demonstrates the most efficient solution: replace NaNs with 0
-    before quantization, which prevents them from contaminating block scales.
-    """
-    set_random_seed(42)
-    device = "cuda:0"
-    torch.set_default_device(device)
-
-    total_tokens = num_tokens + num_padding
-
-    # Create input with NaNs at padding positions (end)
-    x = torch.randn(total_tokens, hidden_size, dtype=dtype, device=device)
-    x[num_tokens:, :] = float('nan')
-
-    # Create mask
-    mask = torch.ones(total_tokens, dtype=torch.bool, device=device)
-    mask[num_tokens:] = False
-
-    # **FIX**: Replace NaNs with 0 before quantization
-    # This is zero-cost if we piggyback on existing attention masking
-    x_masked = torch.where(torch.isnan(x), torch.zeros_like(x), x)
-
-    # Verify masking worked
-    assert not torch.isnan(x_masked).any(), "Masking should remove all NaNs"
-
-    # Create FP4 weights
-    weight_fp4, weight_scale, weight_global_scale, weights_padding_cols = \
-        create_nvfp4_weight(output_size, hidden_size, dtype, device)
-
-    # Compute scales using masked input
-    input_amax = torch.abs(x_masked).max().to(torch.float32)
-    input_global_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / input_amax).to(
-        torch.float32)
-    input_global_scale_inv = 1.0 / input_global_scale
-    alpha = (input_global_scale * weight_global_scale).to(torch.float32)
-
-    # Quantize masked input
-    x_fp4, x_blockscale = ops.scaled_fp4_quant(
-        x_masked, input_global_scale_inv, is_sf_swizzled_layout=True,
-        backend="flashinfer-cutlass")
-
-    # Pad and run GEMM
-    x_fp4_padded = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
-    output = flashinfer_scaled_fp4_mm(
-        x_fp4_padded,
-        weight_fp4,
-        x_blockscale,
-        weight_scale,
-        alpha,
-        dtype,
-        backend="cutlass",
-    )
-    output = slice_nvfp4_output(output, output_size)
-
-    # With the fix, no NaNs should appear in any position
-    assert not torch.isnan(output).any(), (
-        "With NaN masking before quantization, output should be NaN-free"
-    )
-
-    print(f"✓ NaN masking fix verified: no NaNs in output after masking input")
-
-
-@pytest.mark.parametrize("block_size", [16])
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@torch.inference_mode()
-def test_nvfp4_quant_nan_in_block_scale(
-    block_size: int,
-    dtype: torch.dtype,
-) -> None:
-    """
-    Test how NaN affects FP4 block scale computation.
-
-    This isolates the quantization step to understand how NaN in a block
-    affects the block's scaling factor.
-    """
-    device = "cuda:0"
-    torch.set_default_device(device)
-
-    # Create a tensor with one block containing NaN
-    num_blocks = 4
-    x = torch.randn(1, num_blocks * block_size, dtype=dtype, device=device)
-
-    # Inject NaN into the second block
-    x[0, block_size:2*block_size] = float('nan')
-
-    # Compute global scale (will be NaN if computed from the whole tensor)
-    input_amax_with_nan = torch.abs(x).max().to(torch.float32)
-
-    # Compute global scale without NaN (using nanmax equivalent)
-    input_amax_no_nan = torch.abs(x[torch.isfinite(x)]).max().to(torch.float32)
-
-    print(f"Max with NaN: {input_amax_with_nan}")
-    print(f"Max without NaN: {input_amax_no_nan}")
-
-    # If the entire tensor's max is NaN, the global scale is NaN
-    if torch.isnan(input_amax_with_nan):
-        input_global_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX /
-                             input_amax_no_nan).to(torch.float32)
-    else:
-        input_global_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX /
-                             input_amax_with_nan).to(torch.float32)
-
-    input_global_scale_inv = 1.0 / input_global_scale
-
-    # Quantize
-    x_fp4, x_blockscale = ops.scaled_fp4_quant(
-        x, input_global_scale_inv, is_sf_swizzled_layout=False,
-        backend="flashinfer-cutlass")
-
-    # Check block scales
-    print(f"\nBlock scales (FP8): {x_blockscale}")
-    print(f"Block scales (FP32): {x_blockscale.to(torch.float32)}")
-
-    # Check if NaN in one block contaminates neighboring blocks
-    # Convert to float32 for inspection
-    scales_fp32 = x_blockscale.view(torch.float8_e4m3fn).to(torch.float32)
-
-    # The block with NaN will likely have inf or nan scale
-    # Check if this contaminates other blocks
-    has_nan_scale = torch.isnan(scales_fp32).any()
-    has_inf_scale = torch.isinf(scales_fp32).any()
-
-    print(f"Has NaN in block scales: {has_nan_scale}")
-    print(f"Has Inf in block scales: {has_inf_scale}")
-
-    # This test is for observation - we don't assert, just report behavior
-    if has_nan_scale or has_inf_scale:
-        print("⚠ NaN in input produces NaN/Inf in block scales")
-    else:
-        print("✓ Block scales remain finite despite NaN in input")
-
-
-if __name__ == "__main__":
-    # Run a quick smoke test
-    print("Running NVFP4 NaN propagation tests...")
-    test_nvfp4_gemm_nan_isolation(
-        num_tokens=32, num_padding=8, hidden_size=1024, output_size=1024,
-        dtype=torch.bfloat16, nan_placement="end")
-    test_nvfp4_gemm_nan_masking_fix(
-        num_tokens=32, num_padding=8, hidden_size=1024, output_size=1024,
-        dtype=torch.bfloat16)
-    test_nvfp4_quant_nan_in_block_scale(block_size=16, dtype=torch.bfloat16)
-    print("All tests passed!")
@@ -56,11 +56,11 @@ def create_fp4_scale_tensor(
        rounded_m = round_up(m, 128)
        scale_n = n // block_size
        rounded_n = round_up(scale_n, 4)
-        return torch.zeros(
+        return torch.empty(
            (rounded_m, rounded_n // 4), device=device, dtype=torch.int32
        )
    else:
-        return torch.zeros((m, n // block_size), device=device, dtype=torch.uint8)
+        return torch.empty((m, n // block_size), device=device, dtype=torch.uint8)


 def create_fp4_output_tensors(
@@ -403,31 +403,15 @@ def rotary_embedding(

 # layer norm ops
 def rms_norm(
-    out: torch.Tensor,
-    input: torch.Tensor,
-    weight: torch.Tensor,
-    epsilon: float,
-    nan_flags: torch.Tensor | None = None,
-    layer_idx: int = 0,
-    max_num_tokens: int = 0,
+    out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor, epsilon: float
 ) -> None:
-    torch.ops._C.rms_norm(
-        out, input, weight, epsilon, nan_flags, layer_idx, max_num_tokens
-    )
+    torch.ops._C.rms_norm(out, input, weight, epsilon)


 def fused_add_rms_norm(
-    input: torch.Tensor,
-    residual: torch.Tensor,
-    weight: torch.Tensor,
-    epsilon: float,
-    nan_flags: torch.Tensor | None = None,
-    layer_idx: int = 0,
-    max_num_tokens: int = 0,
+    input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, epsilon: float
 ) -> None:
-    torch.ops._C.fused_add_rms_norm(
-        input, residual, weight, epsilon, nan_flags, layer_idx, max_num_tokens
-    )
+    torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)


 def fused_qk_norm_rope(
@@ -73,11 +73,7 @@ def run_rebalance_experts(
    # Move the global expert load window to CPU for computation.
    global_expert_load_window = eplb_stats.global_expert_load_window.cpu()
    # Compute new expert mappings for the model
-    (
-        new_physical_to_logical_map,
-        new_logical_to_physical_map,
-        new_logical_replica_count,
-    ) = eplb_state.policy.rebalance_experts(
+    new_physical_to_logical_map = eplb_state.policy.rebalance_experts(
        global_expert_load_window,
        eplb_stats.num_replicas,
        eplb_stats.num_groups,
@@ -89,16 +85,6 @@ def run_rebalance_experts(

    model_state.new_physical_to_logical_map = new_physical_to_logical_map

-    max_slots = model_state.logical_to_physical_map.shape[-1]
-    padded_logical = torch.nn.functional.pad(
-        new_logical_to_physical_map,
-        (0, max(0, max_slots - new_logical_to_physical_map.shape[-1])),
-        value=-1,
-    ).to(model_state.logical_to_physical_map.device)
-    new_replica = new_logical_replica_count.to(model_state.logical_replica_count.device)
-    model_state.new_logical_to_physical_map = padded_logical
-    model_state.new_logical_replica_count = new_replica
-

 async def transfer_run_periodically(
    state: "EplbState",
@@ -235,16 +235,6 @@ class EplbModelState:
    intermediate variable between `move_to_buffer` and `move_to_workspace`.
    the size is same as physical_to_logical_map
    """
-    new_logical_to_physical_map: torch.Tensor | None = None
-    """
-    intermediate variable between `move_to_buffer` and `move_to_workspace`.
-    the size is same as logical_to_physical_map
-    """
-    new_logical_replica_count: torch.Tensor | None = None
-    """
-    intermediate variable between `move_to_buffer` and `move_to_workspace`.
-    the size is same as logical_replica_count
-    """


 class EplbState:
@@ -508,8 +498,6 @@ class EplbState:
            ),
            cuda_device_index=self.cuda_device_index,
            new_physical_to_logical_map=None,
-            new_logical_to_physical_map=None,
-            new_logical_replica_count=None,
        )
        self.model_states[model_config.compute_hash()] = model_state
        self.num_valid_physical_experts = model.num_physical_experts
@@ -738,17 +726,20 @@ class EplbState:
        ):
            if not self.is_async or is_profile:
                # Get new expert mappings for the model
-                (
-                    new_physical_to_logical_map,
-                    new_logical_to_physical_map,
-                    new_logical_replica_count,
-                ) = self.policy.rebalance_experts(
-                    global_expert_load_window,
+                new_physical_to_logical_map = self.policy.rebalance_experts(
+                    global_expert_load_window.cpu(),
                    num_replicas,
                    num_groups,
                    num_nodes,
                    num_gpus,
-                    eplb_model_state.physical_to_logical_map,
+                    eplb_model_state.physical_to_logical_map.cpu(),
+                )
+
+                num_logical_experts = global_expert_load_window.shape[-1]
+                (new_logical_to_physical_map, new_logical_replica_count) = (
+                    compute_logical_maps(
+                        new_physical_to_logical_map, num_logical_experts
+                    )
                )

                # Update expert weights
@@ -847,11 +838,7 @@ class EplbState:
    def _update_layer_mapping_from_new(
        self, model_state: EplbModelState, layer: int
    ) -> None:
-        if (
-            model_state.new_physical_to_logical_map is None
-            or model_state.new_logical_to_physical_map is None
-            or model_state.new_logical_replica_count is None
-        ):
+        if model_state.new_physical_to_logical_map is None:
            return

        target_device = model_state.physical_to_logical_map.device
@@ -865,19 +852,23 @@ class EplbState:
                new_physical[layer].to(target_device, non_blocking=True)
            )

+        num_logical_experts = model_state.logical_to_physical_map.shape[1]
+        new_logical, new_replica_count = compute_logical_maps(
+            new_physical[layer], num_logical_experts
+        )
+
        logical_device = model_state.logical_to_physical_map.device
-        new_logical = model_state.new_logical_to_physical_map[layer].to(logical_device)
        max_slots = model_state.logical_to_physical_map.shape[-1]
        slot_delta = max_slots - new_logical.shape[-1]
        if slot_delta > 0:
            new_logical = torch.nn.functional.pad(
                new_logical, (0, slot_delta), value=-1
            )
-        model_state.logical_to_physical_map[layer].copy_(new_logical)
+        model_state.logical_to_physical_map[layer].copy_(new_logical.to(logical_device))

        replica_device = model_state.logical_replica_count.device
        model_state.logical_replica_count[layer].copy_(
-            model_state.new_logical_replica_count[layer].to(replica_device)
+            new_replica_count.to(replica_device)
        )

    def _all_ranks_buffer_ready(self, model_state: EplbModelState) -> bool:
@@ -966,7 +957,7 @@ class EplbState:
                transferred_layer,
            )
            if model_state.layer_to_transfer >= model_state.model.num_moe_layers:
-                self.post_eplb(model_state, is_profile)
+                self.post_eplb(model_state)
                model_state.rebalanced = False
                model_state.layer_to_transfer = 0
                model_state.pending_global_ready_check = False
@@ -987,14 +978,9 @@ class EplbState:
                    str(e),
                )

-    def post_eplb(self, model_state: EplbModelState, is_profile: bool = False) -> None:
+    def post_eplb(self, model_state: EplbModelState) -> None:
        assert model_state.new_physical_to_logical_map is not None
-        assert model_state.new_logical_to_physical_map is not None
-        assert model_state.new_logical_replica_count is not None
-
        model_state.new_physical_to_logical_map = None
-        model_state.new_logical_to_physical_map = None
-        model_state.new_logical_replica_count = None

    def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]:
        """
@@ -1052,39 +1038,28 @@ class EplbState:
            model_config=model_config,
        )
        eplb_state.num_valid_physical_experts = num_valid_physical_experts
-        num_moe_layers = expanded_physical_to_logical.shape[0]
-        num_physical_experts = expanded_physical_to_logical.shape[1]
        eplb_model_state = eplb_state.model_states[model_config.compute_hash()]
        eplb_model_state.physical_to_logical_map.copy_(expanded_physical_to_logical)

-        logical_to_physical_map = torch.full(
-            (
-                num_moe_layers,
-                model.num_logical_experts,
-                eplb_model_state.logical_to_physical_map.shape[2],
-            ),
-            -1,
-            dtype=torch.int64,
+        (logical_to_physical_map_cpu, logical_replica_count_cpu) = compute_logical_maps(
+            expanded_physical_to_logical.cpu(), model.num_logical_experts
        )
-        logical_replica_count = torch.zeros(
-            (num_moe_layers, model.num_logical_experts),
-            dtype=torch.int64,
-        )
-        expanded_physical_to_logical_numpy = expanded_physical_to_logical.cpu().numpy()
-        for layer_idx in range(num_moe_layers):
-            for phys_idx in range(num_physical_experts):
-                logical_idx = expanded_physical_to_logical_numpy[layer_idx, phys_idx]
-                if logical_idx >= 0:
-                    replica_idx = logical_replica_count[layer_idx, logical_idx]
-                    logical_to_physical_map[layer_idx, logical_idx, replica_idx] = (
-                        phys_idx
-                    )
-                    logical_replica_count[layer_idx, logical_idx] += 1

-        logical_to_physical_map = logical_to_physical_map.to(device)
-        logical_replica_count = logical_replica_count.to(device)
+        max_num_replicas = eplb_model_state.logical_to_physical_map.shape[-1]
+        num_replicas = logical_to_physical_map_cpu.shape[-1]
+        logical_to_physical_map = torch.nn.functional.pad(
+            logical_to_physical_map_cpu,
+            (
+                0,
+                max_num_replicas - num_replicas,
+            ),
+            value=-1,
+        ).to(device)
+        logical_replica_count = logical_replica_count_cpu.to(device)
+
        eplb_model_state.logical_to_physical_map.copy_(logical_to_physical_map)
        eplb_model_state.logical_replica_count.copy_(logical_replica_count)
+
        return eplb_state


@@ -1132,3 +1107,82 @@ def _node_count_with_rank_mapping(
                node_assignment[other_rank] = next_node_id

    return next_node_id
+
+
+def compute_logical_maps(
+    physical_to_logical_map: torch.Tensor,
+    num_logical_experts: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Derive logical_to_physical_map and logical_replica_count from
+    physical_to_logical_map.
+
+    Args:
+        physical_to_logical_map: [num_layers, num_physical_experts], logical
+            expert index for each physical expert slot
+        num_logical_experts: total number of logical experts
+
+    Returns:
+        logical_to_physical_map: [num_layers, num_logical_experts, max_replicas],
+            physical slots per logical expert; -1 where unused
+        logical_replica_count: [num_layers, num_logical_experts], number of
+            physical replicas per logical expert
+    """
+    device = physical_to_logical_map.device
+    assert physical_to_logical_map.device.type == "cpu"
+
+    dtype = physical_to_logical_map.dtype
+
+    # If computing maps for a single layer, unsqueeze a single element layer dimension
+    per_layer = physical_to_logical_map.dim() == 1
+    physical_to_logical_map_view = physical_to_logical_map
+    if per_layer:
+        physical_to_logical_map_view = physical_to_logical_map.unsqueeze(0)
+    assert len(physical_to_logical_map_view.shape) == 2
+    num_layers, num_physical = physical_to_logical_map_view.shape
+
+    valid_mask = physical_to_logical_map_view >= 0
+    logical_replica_count = torch.zeros(
+        num_layers,
+        num_logical_experts,
+        dtype=dtype,
+        device=device,
+    )
+    logical_replica_count.scatter_add_(
+        1,
+        physical_to_logical_map_view.clamp(min=0),
+        valid_mask.to(dtype),
+    )
+
+    max_replicas = int(logical_replica_count.max().item())
+    logical_to_physical_map_out = torch.full(
+        (num_layers, num_logical_experts, max_replicas),
+        -1,
+        dtype=dtype,
+        device=device,
+    )
+
+    running_count = torch.zeros_like(logical_replica_count)
+    layer_indices = torch.arange(num_layers, device=device)
+    for phys_idx in range(num_physical):
+        # Logical expert at physical slot phys_idx for each layer
+        logical_expert_ids = physical_to_logical_map_view[:, phys_idx]  # [num_layers]
+
+        # Scale up will set the logical expert ids to -1 for all new physical experts.
+        # Only consider "valid" experts when setting up the logical_to_physical map.
+        valid_expert_mask = logical_expert_ids >= 0
+        if not valid_expert_mask.any():
+            continue
+        valid_layers = layer_indices[valid_expert_mask]
+        valid_experts = logical_expert_ids[valid_expert_mask]
+
+        # Use the current running count as the replica index, then increment it.
+        replica_idx = running_count[valid_layers, valid_experts]
+        logical_to_physical_map_out[valid_layers, valid_experts, replica_idx] = phys_idx
+        running_count[valid_layers, valid_experts] += 1
+
+    # If computing maps for a single layer, squeeze out the extra layer dimension
+    # before returning
+    if per_layer:
+        return logical_to_physical_map_out.squeeze(0), logical_replica_count.squeeze(0)
+    return logical_to_physical_map_out, logical_replica_count
@@ -17,7 +17,7 @@ class AbstractEplbPolicy(ABC):
        num_nodes: int,
        num_ranks: int,
        old_global_expert_indices: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> torch.Tensor:
        """
        Entry point for expert-parallelism load balancer.

@@ -35,9 +35,5 @@ class AbstractEplbPolicy(ABC):
        Returns:
            physical_to_logical_map: [layers, num_replicas], the expert
                index of each replica
-            logical_to_physical_map: [layers, num_logical_experts, X],
-                the replica indices for each expert
-            expert_count: [layers, num_logical_experts], number of
-                physical replicas for each logical expert
        """
        raise NotImplementedError
@@ -75,7 +75,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
    @classmethod
    def replicate_experts(
        cls, weight: np.ndarray, num_phy: int
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    ) -> tuple[np.ndarray, np.ndarray]:
        """
        Replicate `num_log` experts to `num_phy` replicas, such that the maximum
        load of all replicas is minimized.
@@ -86,22 +86,19 @@ class DefaultEplbPolicy(AbstractEplbPolicy):

        Returns:
            phy2log: [X, num_phy], logical expert id of each physical expert
-            replica_idx: [X, num_phy], the index of the replica for each logical expert
            logcnt: [X, num_log], number of replicas for each logical expert
        """
        n, num_log = weight.shape
        num_redundant = num_phy - num_log
        assert num_redundant >= 0
        phy2log = np.tile(np.arange(num_phy, dtype=np.int64), (n, 1))
-        replica_idx = np.zeros((n, num_phy), dtype=np.int64)
        logcnt = np.ones((n, num_log), dtype=np.int64)
        arangen = np.arange(n, dtype=np.int64)
        for i in range(num_log, num_phy):
            redundant_indices = np.argmax(weight / logcnt, axis=-1)
            phy2log[:, i] = redundant_indices
-            replica_idx[:, i] = logcnt[arangen, redundant_indices]
            logcnt[arangen, redundant_indices] += 1
-        return phy2log, replica_idx, logcnt
+        return phy2log, logcnt

    @classmethod
    def rebalance_experts_hierarchical(
@@ -111,7 +108,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
        num_groups: int,
        num_nodes: int,
        num_gpus: int,
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    ) -> np.ndarray:
        """
        Parameters:
            weight: [num_moe_layers, num_logical_experts]
@@ -124,10 +121,6 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
        Returns:
            phy2log: [layers, num_replicas], the expert
                index of each replica
-            pphy_replicas_idx: [layers, num_logical_experts, X],
-                the replica indices for each expert
-            logcnt: [layers, num_logical_experts], number of
-                physical replicas for each logical expert
        """
        num_layers, num_logical_experts = weight.shape
        assert num_logical_experts % num_groups == 0
@@ -167,7 +160,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
        tokens_per_mlog = np.take_along_axis(weight, mlog2log, axis=1).reshape(
            -1, num_logical_experts // num_nodes
        )
-        phy2mlog, replicas_idx, mlogcnt = cls.replicate_experts(
+        phy2mlog, mlogcnt = cls.replicate_experts(
            tokens_per_mlog, num_physical_experts // num_nodes
        )

@@ -193,22 +186,15 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
        ).reshape(num_layers, -1)
        # Map node-local logical indices back to global logical expert ids.
        pphy2log = np.take_along_axis(mlog2log, pphy2mlog, axis=1)
-        # Reorder replica ranks to the post-packing physical ordering.
-        pphy_replicas_idx = np.take_along_axis(replicas_idx, pphy2phy, axis=1).reshape(
-            num_layers, -1
-        )
-        # Convert replica counts back to the original logical ordering.
-        logcnt = np.take_along_axis(mlogcnt.reshape(num_layers, -1), log2mlog, axis=1)
-        return pphy2log, pphy_replicas_idx, logcnt
+        return pphy2log

    @classmethod
    def preserve_intragpu_slots(
        cls,
        phy2log: np.ndarray,
-        phy_replicas_idx: np.ndarray,
        num_ranks: int,
        old_phy2log: np.ndarray,
-    ) -> tuple[np.ndarray, np.ndarray]:
+    ) -> np.ndarray:
        """
        Reorder the new mapping per GPU so that experts that remain on the same GPU
        keep their previous slot positions when possible. Incoming experts to that GPU
@@ -218,14 +204,13 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
        """
        num_phy_experts = phy2log.shape[1]
        if num_ranks <= 0 or num_phy_experts % num_ranks != 0:
-            return phy2log, phy_replicas_idx
+            return phy2log

        # Move to CPU and convert to NumPy for processing
        slots_per_gpu = num_phy_experts // num_ranks
        num_layers = phy2log.shape[0]

        post_phy2log = phy2log.copy()
-        post_phy_replicas_idx = phy_replicas_idx.copy()

        for gpu_idx in range(num_ranks):
            start = gpu_idx * slots_per_gpu
@@ -233,7 +218,6 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
            # Experts across all layers for this GPU
            old_local = old_phy2log[:, start:end]  # [layers, slots]
            new_local = phy2log[:, start:end]  # [layers, slots]
-            new_ridx = phy_replicas_idx[:, start:end]  # [layers, slots]

            used_new_indices = np.zeros((num_layers, slots_per_gpu), dtype=bool)
            preserved_positions = np.zeros((num_layers, slots_per_gpu), dtype=bool)
@@ -253,9 +237,6 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
                    post_phy2log[layer_indices, start + slot_idx] = new_local[
                        layer_indices, matched_new_positions
                    ]
-                    post_phy_replicas_idx[layer_indices, start + slot_idx] = new_ridx[
-                        layer_indices, matched_new_positions
-                    ]
                    used_new_indices[layer_indices, matched_new_positions] = True
                    preserved_positions[layer_indices, slot_idx] = True

@@ -287,11 +268,8 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
                    post_phy2log[layer_idx, start + dst_pos] = new_local[
                        layer_idx, src_pos
                    ]
-                    post_phy_replicas_idx[layer_idx, start + dst_pos] = new_ridx[
-                        layer_idx, src_pos
-                    ]

-        return post_phy2log, post_phy_replicas_idx
+        return post_phy2log

    @classmethod
    def rebalance_experts(
@@ -302,7 +280,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
        num_nodes: int,
        num_ranks: int,
        old_global_expert_indices: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> torch.Tensor:
        """
        Entry point for expert-parallelism load balancer.

@@ -321,13 +299,7 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
        Returns:
            phy2log: [layers, num_replicas], the expert
                index of each replica
-            log2phy: [layers, num_logical_experts, X],
-                the replica indices for each expert
-            logcnt: [layers, num_logical_experts], number of
-                physical replicas for each logical expert
        """
-        device = weight.device
-        num_layers, num_logical_experts = weight.shape
        weight_np = weight.float().cpu().numpy()
        old_phy2log_np = (
            old_global_expert_indices.cpu().numpy()
@@ -337,17 +309,13 @@ class DefaultEplbPolicy(AbstractEplbPolicy):

        if num_groups % num_nodes == 0:
            # use hierarchical load-balance policy
-            phy2log_np, phy_replicas_idx_np, logcnt_np = (
-                cls.rebalance_experts_hierarchical(
-                    weight_np, num_replicas, num_groups, num_nodes, num_ranks
-                )
+            phy2log_np = cls.rebalance_experts_hierarchical(
+                weight_np, num_replicas, num_groups, num_nodes, num_ranks
            )
        else:
            # use global load-balance policy
-            phy2log_np, phy_replicas_idx_np, logcnt_np = (
-                cls.rebalance_experts_hierarchical(
-                    weight_np, num_replicas, 1, 1, num_ranks
-                )
+            phy2log_np = cls.rebalance_experts_hierarchical(
+                weight_np, num_replicas, 1, 1, num_ranks
            )

        # Optional postprocessing to preserve slots for experts moving
@@ -355,22 +323,10 @@ class DefaultEplbPolicy(AbstractEplbPolicy):
        # Only apply when the number of GPUs and slots per GPU remain unchanged.
        # Helps to avoid unnecessary weight copying when experts move
        # within the same GPU.
-        if old_global_expert_indices is not None:
-            phy2log_np, phy_replicas_idx_np = cls.preserve_intragpu_slots(
-                phy2log_np, phy_replicas_idx_np, num_ranks, old_phy2log_np
+        if old_phy2log_np is not None:
+            phy2log_np = cls.preserve_intragpu_slots(
+                phy2log_np, num_ranks, old_phy2log_np
            )
-        num_redundant_experts = num_replicas - num_logical_experts
-        maxlogcnt = num_redundant_experts + 1
-        log2phy_np = np.full(
-            (num_layers, num_logical_experts, maxlogcnt), -1, dtype=np.int64
-        )
-        layer_indices = np.arange(num_layers)[:, None]
-        replica_indices = np.tile(
-            np.arange(num_replicas, dtype=np.int64), (num_layers, 1)
-        )
-        log2phy_np[layer_indices, phy2log_np, phy_replicas_idx_np] = replica_indices

-        phy2log = torch.from_numpy(phy2log_np).to(device)
-        log2phy = torch.from_numpy(log2phy_np).to(device)
-        logcnt = torch.from_numpy(logcnt_np).to(device)
-        return phy2log, log2phy, logcnt
+        phy2log = torch.from_numpy(phy2log_np)
+        return phy2log
@@ -1135,7 +1135,6 @@ class NixlConnectorWorker:
        # In progress transfers.
        # [req_id -> list[handle]]
        self._recving_metadata: dict[ReqId, ReqMeta] = {}
-        self._sending_metadata: dict[ReqId, ReqMeta] = {}
        self._recving_transfers = defaultdict[ReqId, list[TransferHandle]](list)
        # Track the expiration time of requests that are waiting to be sent.
        self._reqs_to_send: dict[ReqId, float] = {}
@@ -2225,82 +2224,6 @@ class NixlConnectorWorker:
                            cache, indices, block_size_ratio
                        )

-    @staticmethod
-    def _as_fp8(data: torch.Tensor) -> torch.Tensor:
-        """View uint8 KV cache data as fp8 so torch.isnan works."""
-        if data.dtype == torch.uint8:
-            return data.view(torch.float8_e4m3fn)
-        return data
-
-    def _check_kv_blocks_for_nan(
-        self, req_id: str, block_ids: BlockIds, direction: str
-    ):
-        """Check KV cache blocks for NaN values after transfer.
-
-        Uses a fast two-pass approach: first check all blocks across all
-        layers with a single torch.isnan, then only do the expensive
-        per-layer breakdown if something is found.
-        """
-        all_group_blocks = [g for g in block_ids if len(g) > 0]
-        if not all_group_blocks:
-            return
-
-        # Fast pass: check all layers at once.
-        has_nan = False
-        for cache_or_caches in self.device_kv_caches.values():
-            caches = (
-                [cache_or_caches]
-                if isinstance(cache_or_caches, torch.Tensor)
-                else cache_or_caches
-            )
-            for cache in caches:
-                for group_blocks in all_group_blocks:
-                    indices = torch.tensor(
-                        group_blocks, device=cache.device, dtype=torch.long
-                    )
-                    if torch.isnan(
-                        self._as_fp8(cache[indices])
-                    ).any().item():
-                        has_nan = True
-                        break
-                if has_nan:
-                    break
-            if has_nan:
-                break
-
-        if not has_nan:
-            return
-
-        # Slow pass: per-layer breakdown for diagnosis.
-        for layer_name, cache_or_caches in self.device_kv_caches.items():
-            caches = (
-                [cache_or_caches]
-                if isinstance(cache_or_caches, torch.Tensor)
-                else cache_or_caches
-            )
-            for cache in caches:
-                for group_blocks in all_group_blocks:
-                    indices = torch.tensor(
-                        group_blocks, device=cache.device, dtype=torch.long
-                    )
-                    blocks_data = self._as_fp8(cache[indices])
-                    nan_count = torch.isnan(blocks_data).sum().item()
-                    if nan_count > 0:
-                        total_elements = blocks_data.numel()
-                        logger.error(
-                            "*** NaN DETECTED in KV cache during %s *** "
-                            "req_id=%s, layer=%s, blocks=%s, "
-                            "nan_count=%d, total_elements=%d, "
-                            "nan_pct=%.4f%%",
-                            direction,
-                            req_id,
-                            layer_name,
-                            group_blocks,
-                            nan_count,
-                            total_elements,
-                            100.0 * nan_count / total_elements,
-                        )
-
    def get_finished(self) -> tuple[set[str], set[str]]:
        """
        Get requests that are done sending or recving on this specific worker.
@@ -2333,11 +2256,6 @@ class NixlConnectorWorker:
            if self.use_host_buffer:
                self.sync_recved_kv_to_device(req_id, meta)

-            if envs.VLLM_NIXL_NAN_DETECT:
-                self._check_kv_blocks_for_nan(
-                    req_id, meta.local_physical_block_ids, "recv"
-                )
-
            # post processing for heteroblocksize
            block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(
                meta.remote.engine_id
@@ -2373,7 +2291,6 @@ class NixlConnectorWorker:
            )
            self._reqs_to_process.remove(req_id)
            del self._reqs_to_send[req_id]
-            self._sending_metadata.pop(req_id, None)
            done_sending.add(req_id)

        return done_sending, done_recving
@@ -2421,15 +2338,6 @@ class NixlConnectorWorker:
                    del self.consumer_notification_counts_by_req[req_id]
                    self._reqs_to_process.remove(req_id)
                    self._reqs_to_send.pop(req_id, None)
-
-                    if envs.VLLM_NIXL_NAN_DETECT:
-                        send_meta = self._sending_metadata.pop(req_id, None)
-                        if send_meta is not None:
-                            self._check_kv_blocks_for_nan(
-                                req_id,
-                                send_meta.local_physical_block_ids,
-                                "send",
-                            )
        return notified_req_ids

    def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]:
@@ -2553,14 +2461,6 @@ class NixlConnectorWorker:
            if req_id in self._reqs_to_process:
                self._reqs_to_send[req_id] = expiration_time

-        # Track send-side metadata for NaN detection.
-        if envs.VLLM_NIXL_NAN_DETECT:
-            for req_id, meta in metadata.reqs_to_save.items():
-                meta.local_physical_block_ids = (
-                    self._logical_to_kernel_block_ids(meta.local_block_ids)
-                )
-                self._sending_metadata[req_id] = meta
-
    def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):
        assert meta.remote is not None and self.kv_topo is not None
        remote_ranks = self.kv_topo.get_target_remote_ranks_from_engine_id(
@@ -45,7 +45,6 @@ if TYPE_CHECKING:
    NO_COLOR: bool = False
    VLLM_LOG_STATS_INTERVAL: float = 10.0
    VLLM_TRACE_FUNCTION: int = 0
-    VLLM_NAN_DETECT: bool = False
    VLLM_USE_FLASHINFER_SAMPLER: bool | None = None
    VLLM_PP_LAYER_PARTITION: str | None = None
    VLLM_CPU_KVCACHE_SPACE: int | None = 0
@@ -191,7 +190,6 @@ if TYPE_CHECKING:
    VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16: bool = True
    VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None
    VLLM_NIXL_ABORT_REQUEST_TIMEOUT: int = 480
-    VLLM_NIXL_NAN_DETECT: bool = False
    VLLM_MORIIO_CONNECTOR_READ_MODE: bool = False
    VLLM_MORIIO_QP_PER_TRANSFER: int = 1
    VLLM_MORIIO_POST_BATCH_SIZE: int = -1
@@ -694,10 +692,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
    # If set to 1, vllm will trace function calls
    # Useful for debugging
    "VLLM_TRACE_FUNCTION": lambda: int(os.getenv("VLLM_TRACE_FUNCTION", "0")),
-    # If set to 1, enables zero-overhead NaN/Inf detection in RMSNorm kernels.
-    # Detects per-token NaN/Inf at every layer boundary via the existing
-    # variance reduction. Reports layer names and token positions.
-    "VLLM_NAN_DETECT": lambda: bool(int(os.getenv("VLLM_NAN_DETECT", "0"))),
    # If set, vllm will use flashinfer sampler
    "VLLM_USE_FLASHINFER_SAMPLER": lambda: bool(
        int(os.environ["VLLM_USE_FLASHINFER_SAMPLER"])
@@ -1391,11 +1385,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
    "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": lambda: int(
        os.getenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", "480")
    ),
-    # Enable NaN/Inf detection in KV cache blocks during NIXL transfers.
-    # Logs errors when NaN/Inf values are found on send or receive side.
-    "VLLM_NIXL_NAN_DETECT": lambda: bool(
-        int(os.getenv("VLLM_NIXL_NAN_DETECT", "0"))
-    ),
    # Controls the read mode for the Mori-IO connector
    "VLLM_MORIIO_CONNECTOR_READ_MODE": lambda: (
        os.getenv("VLLM_MORIIO_CONNECTOR_READ_MODE", "False").lower() in ("true", "1")
@@ -53,16 +53,11 @@ def _is_oink_stride_compatible_2d(x_2d: torch.Tensor) -> bool:


 def rms_norm(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    variance_epsilon: float,
-    nan_flags: torch.Tensor | None = None,
-    layer_idx: int = 0,
-    max_num_tokens: int = 0,
+    x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
 ) -> torch.Tensor:
    from vllm import _custom_ops as ops

-    if nan_flags is None and envs.VLLM_BATCH_INVARIANT:
+    if vllm_is_batch_invariant():
        return rms_norm_batch_invariant(x, weight, variance_epsilon)
    out = torch.empty_like(x)
    ops.rms_norm(
@@ -70,9 +65,6 @@ def rms_norm(
        x,
        weight,
        variance_epsilon,
-        nan_flags,
-        layer_idx,
-        max_num_tokens,
    )
    return out

@@ -82,13 +74,10 @@ def fused_add_rms_norm(
    residual: torch.Tensor,
    weight: torch.Tensor,
    variance_epsilon: float,
-    nan_flags: torch.Tensor | None = None,
-    layer_idx: int = 0,
-    max_num_tokens: int = 0,
 ) -> tuple[torch.Tensor, torch.Tensor]:
    from vllm import _custom_ops as ops

-    if nan_flags is None and envs.VLLM_BATCH_INVARIANT:
+    if vllm_is_batch_invariant():
        return rms_norm_batch_invariant(
            x + residual, weight, variance_epsilon
        ), x + residual
@@ -97,9 +86,6 @@ def fused_add_rms_norm(
        residual,
        weight,
        variance_epsilon,
-        nan_flags,
-        layer_idx,
-        max_num_tokens,
    )
    return x, residual

@@ -233,15 +219,6 @@ class RMSNorm(CustomOp):
                    self._use_oink_rmsnorm = False
                    self._use_oink_fused_add_rmsnorm = False

-        # NaN/Inf detection: register this layer with the NaNDetector.
-        self._nan_detect_layer_idx: int = -1
-        if envs.VLLM_NAN_DETECT:
-            from vllm.model_executor.layers.nan_detector import NaNDetector
-
-            self._nan_detect_layer_idx = NaNDetector.get().register(
-                f"RMSNorm_{id(self)}"
-            )
-
    @staticmethod
    def forward_static(
        x: torch.Tensor,
@@ -313,40 +290,6 @@ class RMSNorm(CustomOp):
        if self.variance_size_override is not None:
            return self.forward_native(x, residual)

-        # When NaN detection is enabled, bypass Oink/batch-invariant paths
-        # and use the instrumented vLLM CUDA kernels.
-        nan_flags = None
-        nan_layer_idx = 0
-        nan_max_tokens = 0
-        if getattr(self, "_nan_detect_layer_idx", -1) >= 0:
-            from vllm.model_executor.layers.nan_detector import NaNDetector
-
-            detector = NaNDetector.get()
-            nan_flags = detector.nan_flags
-            nan_layer_idx = self._nan_detect_layer_idx
-            nan_max_tokens = detector.max_num_tokens
-            if nan_flags is not None:
-                add_residual = residual is not None
-                if add_residual:
-                    return fused_add_rms_norm(
-                        x,
-                        residual,
-                        self.weight.data,
-                        self.variance_epsilon,
-                        nan_flags=nan_flags,
-                        layer_idx=nan_layer_idx,
-                        max_num_tokens=nan_max_tokens,
-                    )
-                else:
-                    return rms_norm(
-                        x,
-                        self.weight.data,
-                        self.variance_epsilon,
-                        nan_flags=nan_flags,
-                        layer_idx=nan_layer_idx,
-                        max_num_tokens=nan_max_tokens,
-                    )
-
        # Optional Oink SM100 fast path (no residual). This path is
        # torch.compile-friendly via torch.ops.oink.rmsnorm and preserves
        # 2D layouts (including padded rows) when using the Oink
@@ -1,266 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Zero-overhead NaN/Inf detection via RMSNorm kernel instrumentation
-and pluggable tensor checks."""
-
-from __future__ import annotations
-
-import torch
-import torch.nn as nn
-
-from vllm.logger import init_logger
-
-logger = init_logger(__name__)
-
-
-def _as_fp8(data: torch.Tensor) -> torch.Tensor:
-    """View uint8 KV cache data as fp8 so torch.isnan works."""
-    if data.dtype == torch.uint8:
-        return data.view(torch.float8_e4m3fn)
-    return data
-
-
-class NaNDetector:
-    """Manages per-token NaN/Inf detection flags.
-
-    Singleton. Created lazily when ``VLLM_NAN_DETECT=1``.
-
-    The flag array has shape ``int8[num_checkpoints, max_num_tokens]``.
-    Checkpoints can be:
-
-    * **RMSNorm layers** -- the CUDA kernel writes flags via a pointer
-      argument (zero-cost when disabled).
-    * **Arbitrary tensors** -- call :meth:`check_tensor` which uses
-      ``torch.isfinite`` to check for NaN/Inf.  CUDA-graph compatible.
-
-    Both share the same flag array and reporting path.
-    """
-
-    _instance: NaNDetector | None = None
-
-    def __init__(self) -> None:
-        self._counter: int = 0
-        self._layer_names: dict[int, str] = {}
-        self._max_num_tokens: int = 0
-        self._nan_flags: torch.Tensor | None = None
-        self._host_flags: torch.Tensor | None = None
-        self._finalized: bool = False
-        self._kv_caches: list[torch.Tensor] = []
-
-    @classmethod
-    def get(cls) -> NaNDetector:
-        if cls._instance is None:
-            cls._instance = cls()
-        return cls._instance
-
-    @classmethod
-    def reset(cls) -> None:
-        """Reset singleton (for testing)."""
-        cls._instance = None
-
-    # ------------------------------------------------------------------
-    # Registration (before finalize)
-    # ------------------------------------------------------------------
-
-    def register(self, name: str) -> int:
-        """Register a checkpoint.  Returns its index into the flag array.
-
-        Works for both RMSNorm layers (which pass the index to the CUDA
-        kernel) and arbitrary tensor checks (which pass it to
-        :meth:`check_tensor`).
-        """
-        assert not self._finalized, (
-            "Cannot register new checkpoints after NaNDetector.finalize()"
-        )
-        idx = self._counter
-        self._layer_names[idx] = name
-        self._counter += 1
-        return idx
-
-    # ------------------------------------------------------------------
-    # Properties
-    # ------------------------------------------------------------------
-
-    @property
-    def num_checkpoints(self) -> int:
-        return self._counter
-
-    @property
-    def nan_flags(self) -> torch.Tensor | None:
-        return self._nan_flags
-
-    @property
-    def max_num_tokens(self) -> int:
-        return self._max_num_tokens
-
-    # ------------------------------------------------------------------
-    # Lifecycle
-    # ------------------------------------------------------------------
-
-    def finalize(
-        self,
-        device: torch.device,
-        max_num_tokens: int,
-        kv_caches: list[torch.Tensor] | None = None,
-    ) -> None:
-        """Allocate ``int8[num_checkpoints, max_num_tokens]`` flag tensors."""
-        if self._finalized:
-            return
-        n = self._counter
-        if n == 0:
-            logger.warning(
-                "NaNDetector.finalize() called but nothing registered"
-            )
-            return
-        self._max_num_tokens = max_num_tokens
-        self._nan_flags = torch.zeros(
-            n, max_num_tokens, dtype=torch.int8, device=device
-        )
-        self._host_flags = torch.zeros(
-            n, max_num_tokens, dtype=torch.int8
-        ).pin_memory()
-        if kv_caches is not None:
-            self._kv_caches = kv_caches
-        self._finalized = True
-        logger.info(
-            "NaN/Inf detector initialized: %d checkpoints, max %d tokens "
-            "(%.1f KB flag buffer)",
-            n,
-            max_num_tokens,
-            n * max_num_tokens / 1024,
-        )
-
-    def update_layer_names(self, model: nn.Module) -> None:
-        """Walk model modules to assign readable names."""
-        from vllm.model_executor.layers.layernorm import RMSNorm
-
-        for name, module in model.named_modules():
-            if isinstance(module, RMSNorm) and hasattr(
-                module, "_nan_detect_layer_idx"
-            ):
-                idx = module._nan_detect_layer_idx
-                if idx in self._layer_names:
-                    self._layer_names[idx] = name
-            if hasattr(module, "_nan_detect_indices"):
-                for attr_label, idx in module._nan_detect_indices.items():
-                    if idx in self._layer_names:
-                        self._layer_names[idx] = f"{name}.{attr_label}"
-
-    # ------------------------------------------------------------------
-    # Per-step operations
-    # ------------------------------------------------------------------
-
-    def clear(self) -> None:
-        """Zero flags before each forward pass."""
-        if self._nan_flags is not None:
-            self._nan_flags.zero_()
-
-    def check_tensor(
-        self, tensor: torch.Tensor, checkpoint_idx: int
-    ) -> None:
-        """Check *tensor* for NaN/Inf, writing per-token flags.
-
-        Uses ``torch.isfinite`` -- all ops stay on GPU, no D2H sync.
-        CUDA-graph compatible and fullgraph=True safe.
-
-        FP8 tensors are cast to float16 before checking since
-        ``torch.isfinite`` doesn't support FP8 dtypes.
-
-        Args:
-            tensor: ``[num_tokens, ...]`` tensor to check.
-            checkpoint_idx: index returned by :meth:`register`.
-        """
-        if self._nan_flags is None:
-            return
-        num_tokens = tensor.shape[0]
-        t = tensor.view(num_tokens, -1)
-        if t.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz,
-                       torch.float8_e5m2, torch.float8_e5m2fnuz):
-            t = t.to(torch.float16)
-        has_bad = (~torch.isfinite(t)).any(dim=1)
-        self._nan_flags[checkpoint_idx, :num_tokens].bitwise_or_(
-            has_bad.to(torch.int8)
-        )
-
-    # ------------------------------------------------------------------
-    # Post-forward checking
-    # ------------------------------------------------------------------
-
-    def check(self, num_real_tokens: int) -> None:
-        """D2H copy flags, scan, log results.
-
-        Args:
-            num_real_tokens: Number of real (non-padding) tokens in the
-                current batch.  Flags beyond this index are padding.
-        """
-        if self._nan_flags is None or self._host_flags is None:
-            return
-
-        self._host_flags.copy_(self._nan_flags, non_blocking=False)
-
-        real_flags = self._host_flags[:, :num_real_tokens]
-        pad_flags = self._host_flags[:, num_real_tokens:]
-
-        real_bad = real_flags.any(dim=1).nonzero(as_tuple=True)[0]
-        if len(real_bad) > 0:
-            for layer_idx in real_bad.tolist():
-                token_positions = (
-                    real_flags[layer_idx].nonzero(as_tuple=True)[0].tolist()
-                )
-                name = self._layer_names.get(
-                    layer_idx, f"checkpoint_{layer_idx}"
-                )
-                logger.error(
-                    "NaN/Inf detected in real tokens at '%s' "
-                    "(checkpoint %d), token positions: %s",
-                    name,
-                    layer_idx,
-                    token_positions,
-                )
-
-        pad_bad = pad_flags.any(dim=1).nonzero(as_tuple=True)[0]
-        if len(pad_bad) > 0:
-            logger.debug(
-                "NaN/Inf in padding tokens at %d checkpoints",
-                len(pad_bad),
-            )
-
-        if len(real_bad) > 0:
-            raise RuntimeError(
-                f"NaN/Inf detected at {len(real_bad)} checkpoint(s). "
-                "See ERROR logs above for details."
-            )
-
-    # ------------------------------------------------------------------
-    # KV cache checks (on block assignment)
-    # ------------------------------------------------------------------
-
-    def check_kv_blocks(self, block_ids: list[int]) -> None:
-        """Check specific KV cache blocks for NaN/Inf.
-
-        Called when blocks are recycled from the pool to a new request,
-        to detect stale NaN left by a previous request.
-        """
-        if not block_ids or not self._kv_caches:
-            return
-
-        device = self._kv_caches[0].device
-        indices = torch.tensor(block_ids, device=device, dtype=torch.long)
-
-        for group_idx, kv_cache in enumerate(self._kv_caches):
-            if not isinstance(kv_cache, torch.Tensor):
-                continue
-            if indices.max().item() >= kv_cache.shape[0]:
-                continue
-            blocks = _as_fp8(kv_cache[indices])
-            nan_count = torch.isnan(blocks).sum().item()
-            if nan_count > 0:
-                logger.error(
-                    "Stale NaN in recycled KV cache blocks "
-                    "(group %d, block_ids=%s, nan_count=%d, "
-                    "total_elements=%d)",
-                    group_idx,
-                    block_ids[:10],
-                    nan_count,
-                    blocks.numel(),
-                )
@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Any
 import torch
 from torch.nn.parameter import Parameter

-import vllm.envs as envs
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
@@ -216,18 +216,6 @@ def apply_nvfp4_linear(
    output_dtype = x.dtype
    output_shape = [*x.shape[:-1], output_size]

-    # TODO(performance): This NaN masking adds ~19us/layer overhead (~50% on the
-    # quantization step) because it can't fuse into the CUDA kernel. Proper fixes:
-    # 1. Add NaN check to scaled_fp4_quant CUDA kernel (zero-cost)
-    # 2. Fix upstream attention to never produce NaN
-    # 3. Integrate with check_tensor infrastructure
-    # For now, ~0.6ms overhead for 32 layers is acceptable vs 100% NaN crash.
-    #
-    # Background: NaN in any block causes that block's scale to be NaN, which
-    # contaminates the entire token output during GEMM (100% NaN for that token).
-    # Masking NaN→0 before quantization prevents block scale contamination.
-    x = torch.where(torch.isnan(x), torch.zeros_like(x), x)
-
    # Quantize BF16 or FP16 to (FP4 and interleaved block scale)
    x_fp4, x_blockscale = scaled_fp4_quant(
        x, input_global_scale_inv, is_sf_swizzled_layout=True, backend=backend.value
@@ -27,5 +27,15 @@ def get_layer_params_buffers(layer: torch.nn.Module) -> LayerTensors:


 def get_layer_size(layer: torch.nn.Module) -> int:
-    """Calculate total number of elements across all tensors in a layer."""
-    return sum(tensor.numel() for tensor in get_layer_tensors(layer).values())
+    """Calculate total number of elements across loadable tensors in a layer.
+
+    Excludes SKIP_TENSORS (e.g. _expert_map) which are never moved to meta
+    device and never loaded via weight_loader during layerwise reload.
+    """
+    from .meta import SKIP_TENSORS
+
+    return sum(
+        tensor.numel()
+        for name, tensor in get_layer_tensors(layer).items()
+        if name not in SKIP_TENSORS
+    )
@@ -33,7 +33,6 @@ from torch import nn
 from transformers import DeepseekV2Config, DeepseekV3Config

 import vllm._custom_ops as ops
-import vllm.envs as envs
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, ParallelConfig, VllmConfig, get_current_vllm_config
@@ -535,16 +534,6 @@ class DeepseekV2Attention(nn.Module):
            prefix=f"{prefix}.attn",
        )

-        self.prefix = prefix
-        if envs.VLLM_NAN_DETECT:
-            from vllm.model_executor.layers.nan_detector import NaNDetector
-
-            self._nan_detect_indices = {
-                "attn_output": NaNDetector.get().register(
-                    f"{prefix}.attn_output"
-                ),
-            }
-
    def forward(
        self,
        positions: torch.Tensor,
@@ -588,14 +577,6 @@ class DeepseekV2Attention(nn.Module):
        attn_output = attn_output.view(-1, self.num_local_heads, self.qk_head_dim)[
            ..., : self.v_head_dim
        ].reshape(-1, self.num_local_heads * self.v_head_dim)
-
-        if envs.VLLM_NAN_DETECT and hasattr(self, "_nan_detect_indices"):
-            from vllm.model_executor.layers.nan_detector import NaNDetector
-
-            NaNDetector.get().check_tensor(
-                attn_output, self._nan_detect_indices["attn_output"]
-            )
-
        output, _ = self.o_proj(attn_output)
        return output

@@ -162,6 +162,11 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
        # Share workspace buffer across all executions
        self._workspace = g_sm100_workspace

+        # Pre-allocated output buffer, lazily sized on first call.
+        # Zero-init once to prevent NaN in padding slots (seq_lens=0)
+        # from contaminating downstream per-tensor reductions.
+        self._decode_out: torch.Tensor | None = None
+
    def _sm100_cutlass_mla_decode(
        self,
        q_nope: torch.Tensor,
@@ -218,7 +223,15 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
            if is_quantized_kv_cache(self.kv_cache_dtype)
            else q_nope.dtype
        )
-        out = q_nope.new_empty((B_q, MAX_HEADS, D_latent), dtype=dtype)
+        # Reuse pre-allocated zero-init output buffer to avoid a memset
+        # kernel on every CUDA graph replay.
+        if (
+            self._decode_out is None
+            or self._decode_out.shape[0] < B_q
+            or self._decode_out.dtype != dtype
+        ):
+            self._decode_out = q_nope.new_zeros((B_q, MAX_HEADS, D_latent), dtype=dtype)
+        out = self._decode_out[:B_q]
        lse = (
            torch.empty((B_q, MAX_HEADS), dtype=torch.float32, device=q_nope.device)
            if self.need_to_return_lse_for_decode
@@ -21,6 +21,7 @@ from vllm.v1.attention.backend import (
    AttentionLayer,
    AttentionType,
    MultipleOf,
+    is_quantized_kv_cache,
 )
 from vllm.v1.attention.backends.utils import KVCacheLayoutType

@@ -151,6 +152,11 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
        self.bmm1_scale: float | None = None
        self.bmm2_scale: float | None = None

+        # Pre-allocated output buffer, lazily sized on first call.
+        # Zero-init once to prevent NaN in padding slots (seq_lens=0)
+        # from contaminating downstream per-tensor reductions.
+        self._decode_out: torch.Tensor | None = None
+
    def forward_mqa(
        self,
        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
@@ -186,6 +192,37 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
            if self.kv_cache_dtype.startswith("fp8"):
                self.bmm2_scale *= layer._k_scale_float

+        # Reuse pre-allocated zero-init output buffer to avoid a memset
+        # kernel on every CUDA graph replay.
+        # q is 4D: (batch, q_len_per_req, num_heads, head_dim)
+        # FlashInfer has a bug where out= validation hardcodes 3D shape
+        # (batch, num_heads, kv_lora_rank), but the kernel writes 4D
+        # (batch, q_len, num_heads, kv_lora_rank) when q_len > 1.
+        # So we can only pass out= for single-token decode (q_len == 1).
+        # For q_len > 1, we zero padding slots after the kernel returns.
+        # TODO: upstream fix to FlashInfer
+        B, q_len_per_req = q.shape[0], q.shape[1]
+        out_kwargs: dict[str, torch.Tensor] = {}
+        if q_len_per_req == 1:
+            dtype = (
+                torch.bfloat16
+                if is_quantized_kv_cache(self.kv_cache_dtype)
+                else q.dtype
+            )
+            if (
+                self._decode_out is None
+                or self._decode_out.shape[0] < B
+                or self._decode_out.dtype != dtype
+            ):
+                self._decode_out = torch.zeros(
+                    B,
+                    q.shape[2],
+                    self.kv_lora_rank,
+                    dtype=dtype,
+                    device=q.device,
+                )
+            out_kwargs["out"] = self._decode_out[:B]
+
        o = trtllm_batch_decode_with_kv_cache_mla(
            query=q,
            kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
@@ -198,8 +235,15 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
            max_seq_len=attn_metadata.max_seq_len,
            bmm1_scale=self.bmm1_scale,
            bmm2_scale=self.bmm2_scale,
+            **out_kwargs,
        )

+        # For q_len > 1, we can't pass out= so we work around by zeroing padding slots
+        if not out_kwargs:
+            num_real = attn_metadata.num_decodes
+            if num_real < o.shape[0]:
+                o[num_real:] = 0
+
        # Flatten the output for consistent shape
        o = o.view(-1, o.shape[-2], o.shape[-1])

@@ -883,7 +883,7 @@ class Scheduler(SchedulerInterface):

        new_block_ids_to_zero = (
            (self.kv_cache_manager.take_new_block_ids() or None)
-            if self.needs_kv_cache_zeroing or envs.VLLM_NAN_DETECT
+            if self.needs_kv_cache_zeroing
            else None
        )

@@ -1632,7 +1632,11 @@ class DPEngineCoreProc(EngineCoreProc):
        if self.has_coordinator and request_wave != self.current_wave:
            if request_wave > self.current_wave:
                self.current_wave = request_wave
-            elif not self.engines_running:
+            elif (
+                not self.engines_running
+                and self.scheduler.pause_state == PauseState.UNPAUSED
+            ):
+                self.engines_running = True
                # Request received for an already-completed wave, notify
                # front-end that we need to start the next one.
                self.output_queue.put_nowait(
@@ -413,6 +413,7 @@ class PrometheusStatLogger(AggregateStatLoggerBase):

        labelnames = ["model_name", "engine"]
        model_name = vllm_config.model_config.served_model_name
+        self.model_name = model_name
        max_model_len = vllm_config.model_config.max_model_len

        per_engine_labelvalues: dict[int, list[object]] = {
@@ -975,6 +976,18 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
            self.histogram_kv_block_idle_before_evict = {}
            self.histogram_kv_block_reuse_gap = {}

+        #
+        # CUDAGraph metrics
+        #
+        self._counter_cudagraph_iterations_base = self._counter_cls(
+            name="vllm:cudagraph_iterations",
+            documentation=(
+                "Number of engine iterations by CUDA graph runtime mode."
+            ),
+            labelnames=labelnames + ["runtime_mode"],
+        )
+        self.counter_cudagraph_iterations: dict[str, dict[int, Counter]] = {}
+
        #
        # LoRA metrics
        #
@@ -1086,6 +1099,17 @@ class PrometheusStatLogger(AggregateStatLoggerBase):
                    for gap in event.reuse_gaps_seconds:
                        reuse_hist.observe(gap)

+            if scheduler_stats.cudagraph_stats is not None:
+                mode = scheduler_stats.cudagraph_stats.runtime_mode
+                if mode not in self.counter_cudagraph_iterations:
+                    self.counter_cudagraph_iterations[mode] = {
+                        idx: self._counter_cudagraph_iterations_base.labels(
+                            self.model_name, str(idx), mode
+                        )
+                        for idx in self.engine_indexes
+                    }
+                self.counter_cudagraph_iterations[mode][engine_idx].inc()
+
            if self.gauge_lora_info is not None:
                running_lora_adapters = ",".join(
                    scheduler_stats.running_lora_adapters.keys()
@@ -3,11 +3,14 @@
 import torch

 from vllm.triton_utils import tl, triton
+from vllm.v1.outputs import LogprobsTensors
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.metrics.logits import get_num_nans
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
+from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
 from vllm.v1.worker.gpu.sample.output import SamplerOutput
 from vllm.v1.worker.gpu.sample.sampler import Sampler
+from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS


@triton.jit
@@ -418,6 +421,26 @@ def probabilistic_rejection_sample(
    return sampled, rejected_steps + 1


+@triton.jit
+def _flatten_sampled_kernel(
+    # [num_logits]
+    flat_sampled_ptr,
+    # [num_reqs, num_speculative_steps + 1]
+    sampled_ptr,
+    sampled_stride,
+    # [num_reqs]
+    num_sampled_ptr,
+    # [num_reqs + 1]
+    cu_num_logits_ptr,
+):
+    req_idx = tl.program_id(0)
+    start_idx = tl.load(cu_num_logits_ptr + req_idx)
+    num_sampled = tl.load(num_sampled_ptr + req_idx)
+    for i in range(num_sampled):
+        token_id = tl.load(sampled_ptr + req_idx * sampled_stride + i)
+        tl.store(flat_sampled_ptr + start_idx + i, token_id)
+
+
 class RejectionSampler:
    def __init__(
        self,
@@ -429,6 +452,40 @@ class RejectionSampler:
        self.num_speculative_steps = num_speculative_steps
        self.use_strict_rejection_sampling = use_strict_rejection_sampling

+    def _get_logprobs_tensors(
+        self,
+        input_batch: InputBatch,
+        sampled: torch.Tensor,
+        num_sampled: torch.Tensor,
+        logits: torch.Tensor,
+    ) -> LogprobsTensors | None:
+        max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(
+            input_batch.idx_mapping_np
+        )
+        if max_num_logprobs == NO_LOGPROBS:
+            return None
+
+        num_reqs = input_batch.cu_num_logits.shape[0] - 1
+        num_logits = logits.shape[0]
+        flat_sampled = torch.zeros(
+            num_logits, dtype=sampled.dtype, device=sampled.device
+        )
+        _flatten_sampled_kernel[(num_reqs,)](
+            flat_sampled,
+            sampled,
+            sampled.stride(0),
+            num_sampled,
+            input_batch.cu_num_logits,
+            num_warps=1,
+        )
+        expanded_logits = num_logits != input_batch.idx_mapping.shape[0]
+        return compute_topk_logprobs(
+            logits,
+            max_num_logprobs,
+            flat_sampled,
+            input_batch.cu_num_logits_np.tolist() if expanded_logits else None,
+        )
+
    def __call__(
        self,
        logits: torch.Tensor,
@@ -460,8 +517,6 @@ class RejectionSampler:
                draft_sampled,
                input_batch.expanded_local_pos,
            )
-            # TODO (TheEpicDolphin): Return logprobs for sampled token ids.
-            logprobs_tensors = None
            sampled, num_sampled = probabilistic_rejection_sample(
                processed_logits,
                draft_logits,
@@ -475,6 +530,14 @@ class RejectionSampler:
                self.sampler.sampling_states.seeds.gpu,
                self.num_speculative_steps,
            )
+            logprobs_tensors = self._get_logprobs_tensors(
+                input_batch,
+                sampled,
+                num_sampled,
+                processed_logits
+                if self.sampler.logprobs_mode == "processed_logprobs"
+                else logits,
+            )

        return SamplerOutput(
            sampled_token_ids=sampled,
@@ -1057,26 +1057,7 @@ class GPUModelRunner(

        # Zero GPU memory for freshly allocated cache blocks to prevent
        # stale NaN/data from corrupting attention or SSM computation.
-        #
-        # NOTE: Block zeroing only happens for Mamba/SSM models
-        # (needs_kv_cache_zeroing=True). For standard attention models,
-        # _zero_block_ids is a no-op — recycled blocks keep whatever
-        # the previous request wrote. This means stale NaN from a
-        # previous request can persist in recycled KV cache blocks and
-        # corrupt subsequent requests via attention.
-        #
-        # When VLLM_NAN_DETECT=1, the scheduler also collects new
-        # block IDs (even for attention models) so we can check them
-        # for stale NaN before they are reused.
        if scheduler_output.new_block_ids_to_zero:
-            if envs.VLLM_NAN_DETECT:
-                from vllm.model_executor.layers.nan_detector import (
-                    NaNDetector,
-                )
-
-                NaNDetector.get().check_kv_blocks(
-                    scheduler_output.new_block_ids_to_zero
-                )
            self._zero_block_ids(scheduler_output.new_block_ids_to_zero)

        # Free the cached encoder outputs.
@@ -3450,14 +3431,12 @@ class GPUModelRunner(
                # num_tokens_across_dp will no-longer be valid
                assert batch_descriptor.num_tokens == num_tokens_padded

-        cudagraph_stats = None
-        if self.vllm_config.observability_config.cudagraph_metrics:
-            cudagraph_stats = CUDAGraphStat(
-                num_unpadded_tokens=num_tokens,
-                num_padded_tokens=batch_descriptor.num_tokens,
-                num_paddings=batch_descriptor.num_tokens - num_tokens,
-                runtime_mode=str(cudagraph_mode),
-            )
+        cudagraph_stats = CUDAGraphStat(
+            num_unpadded_tokens=num_tokens,
+            num_padded_tokens=batch_descriptor.num_tokens,
+            num_paddings=batch_descriptor.num_tokens - num_tokens,
+            runtime_mode=str(cudagraph_mode),
+        )

        return (
            cudagraph_mode,
@@ -3810,13 +3789,6 @@ class GPUModelRunner(
        # When spec decode is enabled, defer connector finalization
        # (wait_for_save + clear metadata) until after draft model runs.
        defer_kv_connector_finalize = self.speculative_config is not None
-
-        # Clear NaN/Inf detection flags before the forward pass.
-        if envs.VLLM_NAN_DETECT:
-            from vllm.model_executor.layers.nan_detector import NaNDetector
-
-            NaNDetector.get().clear()
-
        with (
            set_forward_context(
                attn_metadata,
@@ -3843,12 +3815,6 @@ class GPUModelRunner(
                **model_kwargs,
            )

-        # Check NaN/Inf detection flags after the forward pass.
-        if envs.VLLM_NAN_DETECT:
-            from vllm.model_executor.layers.nan_detector import NaNDetector
-
-            NaNDetector.get().check(num_scheduled_tokens)
-
        with record_function_or_nullcontext("gpu_model_runner: postprocess"):
            if self.use_aux_hidden_state_outputs:
                # True when EAGLE 3 is used.
@@ -4653,18 +4619,6 @@ class GPUModelRunner(
            if self.eplb_state.is_async:
                self.eplb_state.start_async_loop()

-        # Initialize NaN/Inf detector if enabled.  Must happen after model
-        # loading (so all RMSNorm layers have registered) but before CUDA graph
-        # capture (so the flag tensor address is stable).
-        if envs.VLLM_NAN_DETECT:
-            from vllm.model_executor.layers.nan_detector import NaNDetector
-
-            detector = NaNDetector.get()
-            detector.update_layer_names(self.model)
-            detector.finalize(
-                self.device, self.max_num_tokens, self.kv_caches
-            )
-
        if (
            self.vllm_config.compilation_config.mode
            == CompilationMode.STOCK_TORCH_COMPILE
Author	SHA1	Message	Date
Tyler Michael SmithandClaude Opus 4.6	119ccde424	[Metrics] Add Prometheus counter for CUDA graph iteration mode Add `vllm:cudagraph_iterations` counter with `runtime_mode` label (NONE/PIECEWISE/FULL) so CUDA graph usage percentage can be computed in PromQL. Also always populate CUDAGraphStat regardless of the `--cudagraph-metrics` flag, since the stat is trivially cheap and that flag should only gate the verbose text log table. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>	2026-03-18 23:11:25 -04:00
Aaron HaoandGitHub	6accb21f2a	[bug] Fix deadlock with pause resume and collective_rpc (#37024 ) Signed-off-by: hao-aaron <ahao@anyscale.com>	2026-03-19 01:49:02 +00:00
Giancarlo DelfinandGitHub	053f3b6309	[Model Runner V2] Spec decode rejection sampler logprobs support (#37237 ) Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>	2026-03-19 01:36:27 +00:00
Aaron HaoandGitHub	5f82706a21	[BUG] Exclude SKIP_TENSORS from get_layer_size() + new weight sync example for dpep (#37334 ) Signed-off-by: ahao-anyscale <ahao@anyscale.com>	2026-03-19 00:45:10 +00:00
Sage Moore GitHub Tyler Michael Smith	c32a58cc2a	[EPLB] Simplify EPLB rearrange by only returning one map (#36267 ) Signed-off-by: Sage Moore <sage@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>	2026-03-18 20:34:00 -04:00