[Bugfix] Fix DSA fusion budgets, RoPE capacity and NUMA startup

.buildkite/test_areas/kernels.yaml

@@ -43,6 +43,8 @@ steps:
   - vllm/kernels/triton/
   - tests/kernels/test_concat_mla_q.py
   - tests/kernels/test_fused_qk_norm_rope_gate.py
+  - vllm/model_executor/layers/rotary_embedding/
+  - vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py
   commands:
     - >-
       pytest -v -s
@@ -70,6 +72,7 @@ steps:
       - vllm/model_executor/layers/fused_moe/router/gate_linear.py
       - vllm/model_executor/layers/fused_moe/router/rocm_fp32_router_gemm.py
       - vllm/model_executor/layers/rotary_embedding/
+      - vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py
       - vllm/_aiter_ops.py
       - vllm/platforms/rocm.py
       - tests/kernels/test_fused_qk_norm_rope_gate.py

tests/compile/test_config.py

@@ -21,6 +21,7 @@
     SchedulerConfig,
     SpeculativeConfig,
     VllmConfig,
+    replace,
 )
 from vllm.config.compilation import CompilationMode, PassConfig
 from vllm.engine.arg_utils import EngineArgs
@@ -282,6 +283,37 @@ def test_enforce_eager_jit_warmup(enable_fault_tolerance):
     assert config.kernel_config.enable_jit_warmup == enable_fault_tolerance
 
 
+@pytest.mark.parametrize("threshold_mb", [0.0, 1.0])
+def test_config_replace_accepts_disabled_allreduce_fusion(threshold_mb):
+    """Autotune cache normalization must preserve valid compile ranges."""
+    config = VllmConfig(
+        model_config=ModelConfig(model="facebook/opt-125m", dtype="bfloat16"),
+        parallel_config=ParallelConfig(
+            tensor_parallel_size=2, distributed_executor_backend="mp"
+        ),
+        scheduler_config=SchedulerConfig.default_factory(
+            max_model_len=2048, max_num_batched_tokens=2048, max_num_seqs=1
+        ),
+        compilation_config=CompilationConfig(
+            pass_config=PassConfig(
+                fuse_allreduce_rms=True,
+                fi_allreduce_fusion_max_size_mb=threshold_mb,
+            )
+        ),
+    )
+    expected_endpoints = set(config.compilation_config.compile_ranges_endpoints or ())
+    rebuilt = replace(config, cache_config=replace(config.cache_config))
+    endpoints = rebuilt.compilation_config.compile_ranges_endpoints
+    assert endpoints is not None
+    assert all(endpoint > 0 for endpoint in endpoints)
+    assert set(endpoints) == expected_endpoints
+    assert endpoints[-1] == config.scheduler_config.max_num_batched_tokens
+    assert (
+        rebuilt.compilation_config.pass_config.fi_allreduce_fusion_max_size_mb
+        == threshold_mb
+    )
+
+
 @pytest.mark.forked
 def test_torch_compile_disable(vllm_runner, monkeypatch):
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")

tests/kernels/core/test_fused_allreduce_gemma_rms_norm.py

@@ -24,6 +24,13 @@
 from vllm.utils.torch_utils import set_random_seed
 
 
+@pytest.mark.parametrize("budget, expected", [(0, None), (8192, 1), (65536, 8)])
+def test_explicit_flashinfer_workspace_budget(budget, expected):
+    from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import _max_token_num
+
+    assert _max_token_num(8, 4096, torch.bfloat16, max_size_bytes=budget) == expected
+
+
 @ensure_current_vllm_config()
 def _worker_fused_ar_norm(
     local_rank,

tests/kernels/core/test_rotary_embedding.py

@@ -9,6 +9,24 @@
 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
 
 
+def test_deepseek_yarn_cache_capacity_preserves_existing_positions(default_vllm_config):
+    from vllm.model_executor.layers.rotary_embedding import get_rope
+
+    parameters = {
+        "rope_type": "deepseek_yarn",
+        "factor": 2,
+        "original_max_position_embeddings": 128,
+        "rope_dim": 64,
+    }
+    with torch.device("cpu"):
+        original = get_rope(128, 256, rope_parameters=parameters, dtype=torch.float32)
+        extended = get_rope(128, 512, rope_parameters=parameters, dtype=torch.float32)
+    assert original.cos_sin_cache.shape == (256, 64)
+    assert extended.cos_sin_cache.shape == (512, 64)
+    assert torch.equal(original.cos_sin_cache, extended.cos_sin_cache[:256])
+    assert torch.isfinite(extended.cos_sin_cache[256:]).all()
+
+
 def rotary_embedding_opcheck(
     rot,
     positions: torch.Tensor,

tests/v1/executor/test_multiproc_executor.py

@@ -110,6 +110,7 @@ def dequeue(*, timeout=None):
         (None, 1, False, False, False),  # nothing to overlap
         (None, 8, True, False, False),  # CPU backend sets OpenMP envs per worker
         (None, 8, False, True, False),  # NUMA wrapper patches the executable
+        (None, 8, False, None, False),  # auto-detection may enable that wrapper
     ],
 )
 def test_can_start_workers_concurrently(

vllm/config/vllm.py

@@ -2703,7 +2703,10 @@ def _set_compile_ranges(self):
                     self.model_config.get_hidden_size()
                     * self.model_config.dtype.itemsize
                 )
-                if compile_range_end is not None and max_token_num < compile_range_end:
+                if (
+                    compile_range_end is not None
+                    and 0 < max_token_num < compile_range_end
+                ):
                     computed_compile_ranges_endpoints.append(max_token_num)
                 else:
                     logger.debug(

vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py

@@ -53,19 +53,36 @@
 _FI_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16)
 
 
-def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None:
+def _max_token_num(
+    tp_size: int,
+    hidden_size: int,
+    dtype: torch.dtype,
+    *,
+    max_size_bytes: int | None = None,
+) -> int | None:
     """Workspace token budget for flashinfer fused all-reduce, or None if the
-    current world size / device is unsupported. Mirrors ``FlashInferAllReduce``."""
+    current world size / device is unsupported. Mirrors ``FlashInferAllReduce``.
+    An explicit zero disables fusion; None retains the platform defaults.
+    """
     from vllm.config.compilation import PassConfig
 
-    max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size)
-    if not max_size_mb:
+    if max_size_bytes is None:
+        max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size)
+        if not max_size_mb:
+            return None
+        max_size_bytes = int(max_size_mb * MiB)
+    if not max_size_bytes:
         return None
     element_size = torch.tensor([], dtype=dtype).element_size()
-    return int(max_size_mb * MiB) // (hidden_size * element_size)
+    return max_size_bytes // (hidden_size * element_size)
 
 
-def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool, int]:
+def _can_use_flashinfer(
+    hidden_states: torch.Tensor,
+    tp_size: int,
+    *,
+    max_size_bytes: int | None = None,
+) -> tuple[bool, int]:
     """Whether the flashinfer fused path applies; returns (ok, max_token_num)."""
     if (
         flashinfer_trtllm_fused_allreduce_norm is None
@@ -82,7 +99,9 @@ def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool
         return False, 0
 
     num_tokens, hidden_size = hidden_states.shape
-    max_token_num = _max_token_num(tp_size, hidden_size, hidden_states.dtype)
+    max_token_num = _max_token_num(
+        tp_size, hidden_size, hidden_states.dtype, max_size_bytes=max_size_bytes
+    )
     if max_token_num is None or num_tokens > max_token_num:
         return False, 0
 

vllm/model_executor/layers/indexer_topk.py

@@ -9,9 +9,47 @@
 from vllm import _custom_ops as ops
 from vllm.config import get_current_vllm_config
 from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
 from vllm.utils.flashinfer import has_flashinfer
 from vllm.v1.worker.workspace import current_workspace_manager
 
+
+@triton.jit
+def _upcast_logits_kernel(src_ptr, dst_ptr, row_ends_ptr, stride, BLOCK: tl.constexpr):
+    row = tl.program_id(0)
+    end = tl.load(row_ends_ptr + row)
+    src_ptr += row.to(tl.int64) * stride
+    dst_ptr += row.to(tl.int64) * stride
+    for start in range(tl.program_id(1) * BLOCK, end, tl.num_programs(1) * BLOCK):
+        offs = start + tl.arange(0, BLOCK)
+        mask = offs < end
+        vals = tl.load(src_ptr + offs, mask=mask)
+        tl.store(dst_ptr + offs, vals.to(tl.float32), mask=mask)
+
+
+def as_fp32_logits(logits: torch.Tensor, row_ends: torch.Tensor) -> torch.Tensor:
+    """Upcast SM100's BF16 MQA logits for the FP32 top-k kernels.
+
+    Only columns below ``row_ends`` (int32, one per row) are converted, the
+    span the consumers read; DeepGEMM's aligned row stride is kept.
+    """
+    if logits.dtype == torch.float32:
+        return logits
+    out = torch.empty_strided(
+        logits.shape, logits.stride(), dtype=torch.float32, device=logits.device
+    )
+    num_rows, width = logits.shape
+    if num_rows == 0 or width == 0:
+        return out
+    assert row_ends.numel() == num_rows and row_ends.dtype == torch.int32
+    block = 1024
+    splits = min(triton.cdiv(width, block), max(1, 1024 // num_rows))
+    _upcast_logits_kernel[(num_rows, splits)](
+        logits, out, row_ends.contiguous(), logits.stride(0), BLOCK=block
+    )
+    return out
+
+
 RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
 
 # cooperative_topk's hard row limit: one cluster wave covers at most this many

vllm/model_executor/layers/rotary_embedding/__init__.py

@@ -308,6 +308,7 @@ def get_rope(
             is_neox_style,
             scaling_factor,
             dtype,
+            cache_max_position=max_position,
             **extra_kwargs,
         )
     elif scaling_type == "longrope":

vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py

@@ -44,8 +44,13 @@ def __init__(
         mscale: float = 1,
         mscale_all_dim: float = 0,
         init_cache: bool = True,
+        cache_max_position: int | None = None,
     ) -> None:
         self.scaling_factor = scaling_factor
+        # Cache capacity does not change YaRN's trained position range.
+        self.cache_max_position = max(
+            max_position_embeddings * scaling_factor, cache_max_position or 0
+        )
         self.beta_fast = beta_fast
         self.beta_slow = beta_slow
         # Get n-d magnitude scaling corrected for interpolation.
@@ -103,7 +108,7 @@ def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
     def _compute_cos_sin_cache(self) -> torch.Tensor:
         inv_freq = self._compute_inv_freq(self.scaling_factor)
         t = torch.arange(
-            self.max_position_embeddings * self.scaling_factor,
+            self.cache_max_position,
             dtype=torch.float32,
         )
         freqs = torch.einsum("i,j -> ij", t, inv_freq)
@@ -244,7 +249,7 @@ def __init__(self, *args, **kwargs):
     def _compute_cos_sin_cache(self) -> torch.Tensor:
         inv_freq = self._compute_inv_freq(self.scaling_factor)
         t = torch.arange(
-            self.max_position_embeddings * self.scaling_factor,
+            self.cache_max_position,
             device=inv_freq.device,
             dtype=torch.float32,
         )

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -33,6 +33,7 @@
 )
 from vllm.model_executor.layers.indexer_topk import (
     RADIX_TOPK_WORKSPACE_SIZE,
+    as_fp32_logits,
     deep_select_topk,
     get_indexer_topk,
     is_deep_select_supported,
@@ -293,42 +294,6 @@ def fused_indexer_q_rope_quant(
     return q_fp8, weights_out
 
 
-@triton.jit
-def _upcast_logits_kernel(src_ptr, dst_ptr, row_ends_ptr, stride, BLOCK: tl.constexpr):
-    row = tl.program_id(0)
-    end = tl.load(row_ends_ptr + row)
-    src_ptr += row.to(tl.int64) * stride
-    dst_ptr += row.to(tl.int64) * stride
-    for start in range(tl.program_id(1) * BLOCK, end, tl.num_programs(1) * BLOCK):
-        offs = start + tl.arange(0, BLOCK)
-        mask = offs < end
-        vals = tl.load(src_ptr + offs, mask=mask)
-        tl.store(dst_ptr + offs, vals.to(tl.float32), mask=mask)
-
-
-def as_fp32_logits(logits: torch.Tensor, row_ends: torch.Tensor) -> torch.Tensor:
-    """Upcast SM100's BF16 MQA logits for the FP32 top-k kernels.
-
-    Only columns below ``row_ends`` (int32, one per row) are converted, the
-    span the consumers read; DeepGEMM's aligned row stride is kept.
-    """
-    if logits.dtype == torch.float32:
-        return logits
-    out = torch.empty_strided(
-        logits.shape, logits.stride(), dtype=torch.float32, device=logits.device
-    )
-    num_rows, width = logits.shape
-    if num_rows == 0 or width == 0:
-        return out
-    assert row_ends.numel() == num_rows and row_ends.dtype == torch.int32
-    block = 1024
-    splits = min(triton.cdiv(width, block), max(1, 1024 // num_rows))
-    _upcast_logits_kernel[(num_rows, splits)](
-        logits, out, row_ends.contiguous(), logits.stride(0), BLOCK=block
-    )
-    return out
-
-
 def _gather_workspace_shapes(
     total_seq_lens: int,
     head_dim: int,

vllm/models/common/ops/fused_allreduce_rms_norm.py

@@ -24,6 +24,8 @@ def fused_allreduce_rms_norm(
     hidden_states: torch.Tensor,
     residual: torch.Tensor,
     norm: RMSNorm,
+    *,
+    max_size_bytes: int | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """All-reduce + add residual + (standard) RMSNorm, fused via flashinfer.
 
@@ -32,13 +34,16 @@ def fused_allreduce_rms_norm(
     after. Returns ``(normed_output, new_residual)``, equivalent to
     ``norm(all_reduce(hidden_states), residual)``. Falls back to an explicit
     all-reduce + RMSNorm when the flashinfer fast path is unavailable.
+    ``max_size_bytes`` overrides the default budget; zero disables fusion.
     """
     tp_size = get_tensor_model_parallel_world_size()
     if tp_size == 1:
         return norm(hidden_states, residual)
 
     if flashinfer_trtllm_fused_allreduce_norm is not None:
-        ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size)
+        ok, max_token_num = _can_use_flashinfer(
+            hidden_states, tp_size, max_size_bytes=max_size_bytes
+        )
         if ok:
             norm_out = torch.empty_like(hidden_states)
             # With norm_out provided, the kernel writes the new residual

vllm/models/deepseek_v32/attention.py

@@ -148,7 +148,10 @@ def __init__(
         num_local_heads = num_heads // tp_size
         qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
         scaling = qk_head_dim**-0.5
-        max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
+        max_position_embeddings = max(
+            getattr(config, "max_position_embeddings", 8192),
+            vllm_config.model_config.max_model_len,
+        )
 
         # DSA checkpoints may use plain ("default") or yarn-scaled RoPE.
         if config.rope_parameters["rope_type"] != "default":