diff --git a/docs/contributing/jit_kernel_warmup.md b/docs/contributing/jit_kernel_warmup.md
index 4cf0212790..4377b40427 100644
--- a/docs/contributing/jit_kernel_warmup.md
+++ b/docs/contributing/jit_kernel_warmup.md
@@ -17,7 +17,7 @@ Each warmable kernel defines its compile-key mapping and compile-only entry poin
 
 Here, a **kernel wrapper** (or just **wrapper**) is an instance of a concrete `VllmJitKernel` subclass.
 
-Expose one wrapper near the kernel's normal runtime entry point. Prefer this shape:
+Expose one wrapper near the kernel's normal runtime entry point. Prefer this backend-agnostic shape:
 
 ```python
 class MyKernel(VllmJitKernel["MyKernel.CompileKey"]):
@@ -43,14 +43,14 @@ class MyKernel(VllmJitKernel["MyKernel.CompileKey"]):
         return self.kernel(...)
 
 
-MY_KERNEL = MyKernel()
+_MY_KERNEL = MyKernel()
 ```
 
 `CompileKey`, `dispatch(...)`, and `get_warmup_keys(...)` are backend-agnostic. Backend-specific behavior belongs in `kernel(...)`, `compile(...)`, and `__call__(...)`.
 
 The module-level singleton should be used by warmup and by the runtime call path. This keeps dispatch behavior shared instead of duplicated.
 
-`VllmJitKernel.warmup(...)` compiles every key returned by `get_warmup_keys(...)`; wrappers should not reimplement it.
+`VllmJitKernel.warmup(...)` materializes every key returned by `get_warmup_keys(...)`; wrappers should not reimplement it.
 
 ### Choose Compile-Key Fields
 
@@ -70,18 +70,20 @@ def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
 
 Use independent ranges or alternatives for cartesian products, `zip_inputs(...)` for coupled rows, and `_when` for validity constraints. The complete syntax is documented in [Search-Space Reference](#2-search-space-reference).
 
-### Compile Without Launching
+### Materialize Without Launching
 
-`compile(compile_key)` means "make this specialization available". Depending on the backend, that may compile from source, call a compile-only API, load an already-built artifact, or compile on cache miss.
+The base contract uses `compile(compile_key)` internally to mean "make this specialization available". Depending on the backend, that may compile from source, call a compile-only API, load an already-built artifact, or compile on cache miss. Contributors only implement this hook when their backend helper does not provide it.
 
-`compile(...)` should not launch a real inference workload or allocate real tensors. Each DSL should expose fake tensor/spec descriptors suitable for compilation only.
+For Triton, implement `warmup_inputs(...)`. `@kernel_launcher` binds those inputs to `__call__(...)`, whose body returns the runtime grid and derived launch arguments as a `LaunchSpec`. The shared helper then routes that specification to the normal kernel launch or to `kernel.warmup(...)` during startup.
+
+Materialization should not launch a real inference workload or allocate real tensors. `TritonWarmupTensor(..., strides=None)` represents compact row-major storage; provide explicit representative strides for padded, transposed, or otherwise strided runtime inputs.
 
 ### Register the Selected Wrapper
 
 Register the wrapper where the runtime implementation is selected:
 
 ```python
-MY_KERNEL.register_warmup()
+_MY_KERNEL.register_warmup()
 ```
 
 Registration records metadata only. It does not compile or launch the kernel. Repeated registrations from equivalent layers are allowed and deduplicated later.
@@ -92,7 +94,7 @@ Registration records metadata only. It does not compile or launch the kernel. Re
 - Keep specialization mapping in `dispatch(...)` instead of duplicating it in warmup code.
 - Use fake tensors or backend compile-only descriptors; never perform a dummy runtime launch.
 - Keep registration metadata-only so model construction remains cheap and side-effect free.
-- Compile registered kernels only through `kernel_warmup()`.
+- Materialize registered kernels only through `kernel_warmup()`.
 - Keep runtime execution and startup compilation separate and easy to review.
 - Use one module-level wrapper instance for registration and runtime calls.
 
diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py
index 4b4e10df58..fabdb5bb5b 100644
--- a/tests/model_executor/layers/test_fused_shared_expert.py
+++ b/tests/model_executor/layers/test_fused_shared_expert.py
@@ -553,6 +553,7 @@ def test_models_fse_init(
                     method="mtp",
                     parallel_drafting=False,
                     enable_adaptive_verification=False,
+                    use_dspark=lambda: False,
                 )
                 mtp = DeepSeekV4MTP(vllm_config=vllm_config)
         assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude)
diff --git a/tests/model_executor/test_jit_warmup.py b/tests/model_executor/test_jit_warmup.py
index b5c6f90ea4..29495a510a 100644
--- a/tests/model_executor/test_jit_warmup.py
+++ b/tests/model_executor/test_jit_warmup.py
@@ -590,6 +590,35 @@ def test_registry_passes_vllm_config_to_default_requests() -> None:
     assert kernel.compiled == [ConfigKernel.CompileKey(value=7)]
 
 
+def test_registry_default_request_supports_config_free_kernel() -> None:
+    class ConfigFreeKernel(VllmJitKernel["ConfigFreeKernel.CompileKey"]):
+        @dataclass(frozen=True)
+        class CompileKey:
+            value: int
+
+        def __init__(self) -> None:
+            self.compiled: list[ConfigFreeKernel.CompileKey] = []
+            super().__init__()
+
+        def dispatch(self, *, value: int) -> CompileKey:  # type: ignore[override]
+            return self.CompileKey(value=value)
+
+        def get_warmup_keys(self) -> list[CompileKey]:
+            return [self.dispatch(value=7)]
+
+        def compile(self, compile_key: CompileKey) -> None:
+            self.compiled.append(compile_key)
+
+    registry = JitWarmupRegistry(_config())
+    kernel = ConfigFreeKernel()
+
+    with registry.activate():
+        kernel.register_warmup()
+    registry.warmup()
+
+    assert kernel.compiled == [ConfigFreeKernel.CompileKey(value=7)]
+
+
 def test_get_ast_full_name_handles_names_attributes_and_other_nodes() -> None:
     dotted_expr = ast.parse("foo.bar.baz").body[0]
     call_expr = ast.parse("foo()").body[0]
diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
index 06625e29e2..c0ce30d214 100644
--- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
@@ -8,10 +8,16 @@ import torch
 
 from tests.v1.attention.utils import create_vllm_config
 from vllm.v1.attention.backend import CommonAttentionMetadata
+from vllm.v1.attention.backends.mla.compressor_utils import (
+    CompressedSlotMappingKernel,
+)
 from vllm.v1.attention.backends.mla.indexer import (
     BuildPrefillChunkMetadataKernel,
     DeepseekV32IndexerMetadataBuilder,
 )
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    ConvertReqIndexToGlobalIndexKernel,
+)
 from vllm.v1.kv_cache_interface import MLAAttentionSpec
 from vllm.v1.worker.block_table import get_block_table_width
 
@@ -20,7 +26,7 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():
     config = SimpleNamespace(
         scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
         model_config=SimpleNamespace(
-            hf_config=SimpleNamespace(compress_ratios=[0, 0, 4, 128, 0])
+            hf_config=SimpleNamespace(compress_ratios=[0, 0, 4, 128, 0], index_kpool=32)
         ),
         parallel_config=SimpleNamespace(
             decode_context_parallel_size=1,
@@ -30,7 +36,42 @@ def test_indexer_warmup_normalizes_zero_compress_ratios():
 
     keys = BuildPrefillChunkMetadataKernel().get_warmup_keys(config)
 
-    assert {key.COMPRESS_RATIO for key in keys} == {1, 4, 128}
+    assert {key.compress_ratio for key in keys} == {1, 4, 32, 128}
+    assert {(key.query_slice_start, key.query_slice_stop) for key in keys} == {
+        (query_slice_start, query_slice_stop)
+        for query_slice_start in (1, 2, 16)
+        for query_slice_stop in (1, 2, 16)
+    }
+
+
+def test_compressed_slot_mapping_warmup_includes_index_kpool():
+    config = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=256),
+        model_config=SimpleNamespace(hf_config=SimpleNamespace(index_kpool=32)),
+    )
+
+    keys = CompressedSlotMappingKernel().get_warmup_keys(config)
+    assert {(key.compress_ratio, key.block_size) for key in keys} == {(32, 2)}
+
+
+def test_index_conversion_warmup_uses_physical_block_stride():
+    config = SimpleNamespace(
+        cache_config=SimpleNamespace(block_size=64),
+        model_config=SimpleNamespace(
+            max_model_len=1024,
+            hf_config=SimpleNamespace(index_topk=2048),
+        ),
+        parallel_config=SimpleNamespace(
+            decode_context_parallel_size=1,
+            cp_kv_cache_interleave_size=1,
+        ),
+    )
+
+    keys = ConvertReqIndexToGlobalIndexKernel().get_warmup_keys(
+        config,
+        block_stride_rows=4096,
+    )
+    assert {key.block_stride_rows for key in keys} == {4096}
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py
index 356b313515..559af5572d 100644
--- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py
@@ -7,6 +7,7 @@ from typing import Any
 
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
 from vllm.v1.worker.gpu import eplb_utils as eplb
 from vllm.v1.worker.gpu import model_runner as mrv2
@@ -84,6 +85,7 @@ def _make_runner(**overrides: Any) -> Any:
         post_forward=lambda *_, **__: None,
     )
     runner.eplb = eplb.EPLBController(runner.parallel_config, runner.device)
+    runner.jit_warmup_registry = JitWarmupRegistry(runner.vllm_config)
     runner.pooling_runner = None
     runner.execute_model_state = None
     for key, value in overrides.items():
diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py
index 305b402858..9f06b64177 100644
--- a/vllm/model_executor/layers/attention/mla_attention.py
+++ b/vllm/model_executor/layers/attention/mla_attention.py
@@ -604,6 +604,33 @@ class MLAAttention(nn.Module, AttentionLayerBase):
 
         self.use_sparse = use_sparse
 
+        if vllm_config.kernel_config.enable_jit_warmup:
+            backend_name = self.attn_backend.get_name()
+            if backend_name in (
+                "FLASHMLA_SPARSE",
+                "FLASHINFER_MLA_SPARSE",
+                "FLASHINFER_MLA_SPARSE_SM120",
+                "DEEPSEEK_V32_INDEXER",
+            ):
+                from vllm.v1.attention.backends.mla.compressor_utils import (
+                    _COMPRESSED_SLOT_MAPPING_KERNEL,
+                )
+                from vllm.v1.attention.backends.mla.indexer import (
+                    _BUILD_PREFILL_CHUNK_METADATA_KERNEL,
+                    _PREPARE_UNIFORM_DECODE_KERNEL,
+                )
+
+                _COMPRESSED_SLOT_MAPPING_KERNEL.register_warmup()
+                _PREPARE_UNIFORM_DECODE_KERNEL.register_warmup()
+                _BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup()
+
+                if backend_name != "DEEPSEEK_V32_INDEXER":
+                    from vllm.v1.attention.backends.mla.sparse_swa import (
+                        _COMPUTE_PREFILL_METADATA_KERNEL,
+                    )
+
+                    _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()
+
         self.dcp_manager: MLADCPManager | None = None
         if self.impl.dcp_world_size > 1:
             query_dtype = (
@@ -652,6 +679,26 @@ class MLAAttention(nn.Module, AttentionLayerBase):
     def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
         # [B, H=1, N, C] -> [B, N, C]
         self.kv_cache = kv_cache.squeeze(1)
+        if (
+            self._vllm_config.kernel_config.enable_jit_warmup
+            and self.attn_backend.get_name()
+            in (
+                "FLASHMLA_SPARSE",
+                "FLASHINFER_MLA_SPARSE",
+                "FLASHINFER_MLA_SPARSE_SM120",
+                "DEEPSEEK_V32_INDEXER",
+            )
+        ):
+            from vllm.v1.attention.backends.mla.sparse_utils import (
+                _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL,
+            )
+
+            row_width = self.kv_cache.shape[-1]
+            assert self.kv_cache.stride(0) % row_width == 0
+            _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL.register_warmup(
+                self._vllm_config,
+                block_stride_rows=self.kv_cache.stride(0) // row_width,
+            )
 
     @property
     def chunked_prefill_workspace_size(self) -> int:
diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py
index 9c43bbfc33..99c387a3d1 100644
--- a/vllm/model_executor/warmup/jit_warmup.py
+++ b/vllm/model_executor/warmup/jit_warmup.py
@@ -737,7 +737,11 @@ class JitWarmupRegistry:
         for kernel, registrations in self._registrations.items():
             compile_keys: dict[Any, None] = {}
             for args, kwargs in registrations:
-                if not args and not kwargs:
+                if (
+                    not args
+                    and not kwargs
+                    and inspect.signature(kernel.get_warmup_keys).parameters
+                ):
                     args = (self.vllm_config,)
                 for compile_key in kernel.get_warmup_keys(*args, **kwargs):
                     compile_keys[compile_key] = None
diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py
index 29da528c93..3a87f33599 100644
--- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py
+++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py
@@ -2,15 +2,21 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import ast
 import inspect
-from collections.abc import Callable
+from abc import abstractmethod
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass
-from typing import Any
+from functools import cached_property, wraps
+from typing import Any, ClassVar, Generic, TypeVar, cast
 
 from vllm.model_executor.warmup.jit_warmup import (
+    VllmJitKernel,
     get_ast_full_name,
     get_function_source_node,
 )
 
+CompileKeyT = TypeVar("CompileKeyT")
+LaunchSpec = tuple[tuple[int, ...], dict[str, Any]]
+
 
 def triton_scalar_specialization_rep(value: int) -> int:
     """Return an integer with the same default Triton JIT specialization.
@@ -50,10 +56,16 @@ def triton_scalar_specialization_rep(value: int) -> int:
 
 @dataclass(frozen=True)
 class TritonWarmupTensor:
-    # Compile-only tensor descriptor for Triton pointer specialization.
+    """Compile-only tensor metadata used by Triton warmup.
+
+    ``strides=None`` represents compact row-major storage. Pass explicit strides
+    whenever the runtime tensor can be padded, transposed, or otherwise strided.
+    """
+
     dtype: Any
     aligned: bool = True
     shape: tuple[int, ...] = (1,)
+    strides: tuple[int, ...] | None = None
 
     def data_ptr(self) -> int:
         return 0 if self.aligned else 1
@@ -61,13 +73,81 @@ class TritonWarmupTensor:
     def ptr_range(self) -> int:
         return 0
 
-    def stride(self) -> tuple[int, ...]:
-        strides: list[int] = []
-        stride = 1
-        for size in reversed(self.shape):
-            strides.append(stride)
-            stride *= size
-        return tuple(reversed(strides))
+    def stride(self, dim: int | None = None) -> int | tuple[int, ...]:
+        if self.strides is None:
+            strides: list[int] = []
+            stride = 1
+            for size in reversed(self.shape):
+                strides.append(stride)
+                stride *= size
+            result = tuple(reversed(strides))
+        else:
+            result = self.strides
+        return result if dim is None else result[dim]
+
+
+class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
+    """Triton owner whose runtime launch specification is reused for warmup."""
+
+    kernel: ClassVar[Any]
+    _warming = False
+
+    @abstractmethod
+    def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:
+        """Return runtime-shaped inputs that reproduce one compile key."""
+        raise NotImplementedError
+
+    def compile(self, compile_key: CompileKeyT) -> None:
+        inputs = self.warmup_inputs(compile_key)
+        self._warming = True
+        try:
+            cast(Callable[..., None], self)(**inputs)
+        finally:
+            self._warming = False
+
+    @cached_property
+    def _kernel_param_names(self) -> frozenset[str]:
+        return frozenset(self.kernel.arg_names)
+
+    def launch(
+        self,
+        grid: tuple[int, ...],
+        inputs: Mapping[str, Any],
+        /,
+        **kwargs: Any,
+    ) -> Any:
+        for name, value in inputs.items():
+            target = name if name in self._kernel_param_names else f"{name}_ptr"
+            if target in self._kernel_param_names and target not in kwargs:
+                kwargs[target] = value
+        if self._warming:
+            warmup = getattr(self.kernel, "warmup", None)
+            assert warmup is not None
+            return warmup(grid=(1,), **kwargs)
+        return self.kernel[grid](**kwargs)
+
+
+def kernel_launcher(
+    call_fn: Callable[..., LaunchSpec],
+) -> Callable[..., None]:
+    """Launch a Triton kernel from a declarative ``__call__`` specification."""
+    signature = inspect.signature(call_fn)
+
+    @wraps(call_fn)
+    def wrapper(
+        self: VllmTritonJitKernel[Any],
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        grid, launch_kwargs = call_fn(self, *args, **kwargs)
+        bound = signature.bind(self, *args, **kwargs)
+        bound.apply_defaults()
+        inputs = {
+            name: value for name, value in bound.arguments.items() if name != "self"
+        }
+        self.launch(grid, inputs, **launch_kwargs)
+
+    return wrapper
 
 
 @dataclass(frozen=True)
diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py
index f4dc088447..42b3746e51 100644
--- a/vllm/model_executor/warmup/kernel_warmup.py
+++ b/vllm/model_executor/warmup/kernel_warmup.py
@@ -35,9 +35,6 @@ from vllm.model_executor.warmup.kimi_k3_triton_warmup import (
     kimi_k3_triton_warmup,
 )
 from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
-from vllm.model_executor.warmup.sparse_mla_triton_warmup import (
-    sparse_mla_triton_warmup,
-)
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import is_deep_gemm_supported
 from vllm.utils.flashinfer import has_flashinfer
@@ -154,7 +151,6 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
     if worker.vllm_config.kernel_config.enable_jit_warmup:
         kimi_k3_triton_warmup(worker)
         fa4_cutedsl_warmup(worker)
-        sparse_mla_triton_warmup(worker)
 
     if current_platform.has_device_capability(90):
         _warmup_ll_bf16_router_gemm(worker.get_model())
diff --git a/vllm/model_executor/warmup/sparse_mla_triton_warmup.py b/vllm/model_executor/warmup/sparse_mla_triton_warmup.py
deleted file mode 100644
index cf35bf40e8..0000000000
--- a/vllm/model_executor/warmup/sparse_mla_triton_warmup.py
+++ /dev/null
@@ -1,111 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Warm up sparse-MLA Triton metadata kernels."""
-
-from typing import TYPE_CHECKING
-
-from vllm.logger import init_logger
-
-if TYPE_CHECKING:
-    from vllm.config import VllmConfig
-    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
-    from vllm.v1.worker.gpu_worker import Worker
-
-logger = init_logger(__name__)
-
-_DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset(
-    {
-        "FLASHMLA_SPARSE_DSV4",
-        "FLASHINFER_MLA_SPARSE_DSV4",
-        "ROCM_FLASHMLA_SPARSE_DSV4",
-        "DEEPSEEK_SPARSE_SWA",
-    }
-)
-_GENERIC_SPARSE_MLA_BACKENDS = frozenset(
-    {
-        "FLASHMLA_SPARSE",
-        "FLASHINFER_MLA_SPARSE",
-        "FLASHINFER_MLA_SPARSE_SM120",
-    }
-)
-_INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"})
-
-_INDEXER_PREFILL_CHUNK_METADATA_BACKENDS = frozenset({"DEEPSEEK_V32_INDEXER"})
-
-
-def _attention_backend_name(backend: object) -> str | None:
-    get_name = getattr(backend, "get_name", None)
-    if get_name is None:
-        return None
-    try:
-        return get_name()
-    except NotImplementedError:
-        return None
-
-
-def _has_attention_backend(
-    runner: "GPUModelRunner",
-    backend_names: frozenset[str],
-) -> bool:
-    for groups in getattr(runner, "attn_groups", []) or ():
-        for group in groups:
-            name = _attention_backend_name(getattr(group, "backend", None))
-            if name in backend_names:
-                return True
-    return False
-
-
-def _compile_sparse_swa_prefill_metadata_kernel(
-    vllm_config: "VllmConfig",
-) -> None:
-    from vllm.v1.attention.backends.mla.sparse_swa import (
-        _COMPUTE_PREFILL_METADATA_KERNEL,
-    )
-
-    _COMPUTE_PREFILL_METADATA_KERNEL.warmup(vllm_config)
-
-
-def _compile_prefill_chunk_metadata_kernel(
-    vllm_config: "VllmConfig",
-) -> None:
-    from vllm.v1.attention.backends.mla.indexer import (
-        _BUILD_PREFILL_CHUNK_METADATA_KERNEL,
-    )
-
-    _BUILD_PREFILL_CHUNK_METADATA_KERNEL.warmup(vllm_config)
-
-
-def _compile_combine_topk_swa_indices_kernel(
-    vllm_config: "VllmConfig",
-) -> None:
-    from vllm.models.deepseek_v4.common.ops.cache_utils import (
-        _COMBINE_TOPK_SWA_INDICES_KERNEL,
-    )
-
-    _COMBINE_TOPK_SWA_INDICES_KERNEL.warmup(vllm_config)
-
-
-def sparse_mla_triton_warmup(worker: "Worker") -> None:
-    runner = worker.model_runner
-    if runner.is_pooling_model:
-        return
-
-    max_tokens = worker.scheduler_config.max_num_batched_tokens
-    max_num_prefills = min(worker.scheduler_config.max_num_seqs, max_tokens)
-    if max_tokens <= 0 or max_num_prefills <= 0:
-        return
-
-    vllm_config = runner.vllm_config
-    try:
-        if _has_attention_backend(runner, _DEEPSEEK_V4_SPARSE_MLA_BACKENDS):
-            _compile_sparse_swa_prefill_metadata_kernel(vllm_config)
-            _compile_prefill_chunk_metadata_kernel(vllm_config)
-            _compile_combine_topk_swa_indices_kernel(vllm_config)
-        elif _has_attention_backend(runner, _GENERIC_SPARSE_MLA_BACKENDS):
-            _compile_sparse_swa_prefill_metadata_kernel(vllm_config)
-            _compile_prefill_chunk_metadata_kernel(vllm_config)
-        elif _has_attention_backend(runner, _INDEXER_PREFILL_CHUNK_METADATA_BACKENDS):
-            _compile_prefill_chunk_metadata_kernel(vllm_config)
-
-    except Exception:
-        logger.warning("Skipping sparse MLA Triton warmup.", exc_info=True)
diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py
index 61bbda38c6..e53cc7c2a8 100644
--- a/vllm/models/deepseek_v4/attention.py
+++ b/vllm/models/deepseek_v4/attention.py
@@ -353,6 +353,57 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
                 k_cache_prefix=self.prefix,
             )
 
+        if vllm_config.kernel_config.enable_jit_warmup:
+            from vllm.v1.attention.backends.mla.sparse_swa import (
+                _COMPUTE_PREFILL_METADATA_KERNEL,
+                _COMPUTE_SWA_INDICES_AND_LENS_KERNEL,
+            )
+
+            _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup()
+            _COMPUTE_SWA_INDICES_AND_LENS_KERNEL.register_warmup(
+                window_size=self.window_size,
+                block_size=self.swa_cache_layer.block_size,
+            )
+
+            if self.compress_ratio > 1:
+                from vllm.v1.attention.backends.mla.compressor_utils import (
+                    _COMPRESSED_SLOT_MAPPING_KERNEL,
+                )
+
+                _COMPRESSED_SLOT_MAPPING_KERNEL.register_warmup()
+
+            if self.indexer is not None:
+                from vllm.v1.attention.backends.mla.indexer import (
+                    _BUILD_PREFILL_CHUNK_METADATA_KERNEL,
+                    _PREPARE_UNIFORM_DECODE_KERNEL,
+                )
+
+                _PREPARE_UNIFORM_DECODE_KERNEL.register_warmup()
+                _BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup()
+
+            spec_config = vllm_config.speculative_config
+            if spec_config is not None and spec_config.use_dspark():
+                from vllm.v1.attention.backends.mla.sparse_swa import (
+                    _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL,
+                )
+
+                _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL.register_warmup(
+                    window_size=self.window_size,
+                    num_speculative_tokens=spec_config.num_speculative_tokens,
+                    block_size=self.swa_cache_layer.block_size,
+                )
+
+            if self.backend_cls.get_name() in (
+                "FLASHMLA_SPARSE_DSV4",
+                "ROCM_FLASHMLA_SPARSE_DSV4",
+                "XPU_V4_MLA_SPARSE",
+            ):
+                from vllm.models.deepseek_v4.common.ops.cache_utils import (
+                    _COMBINE_TOPK_SWA_INDICES_KERNEL,
+                )
+
+                _COMBINE_TOPK_SWA_INDICES_KERNEL.register_warmup()
+
     def forward(
         self,
         positions: torch.Tensor,
diff --git a/vllm/v1/attention/backends/mla/compressor_utils.py b/vllm/v1/attention/backends/mla/compressor_utils.py
index 47d331ba7b..a7d801b040 100644
--- a/vllm/v1/attention/backends/mla/compressor_utils.py
+++ b/vllm/v1/attention/backends/mla/compressor_utils.py
@@ -1,7 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+from typing import Any
+
 import torch
 
+from vllm.model_executor.warmup.jit_warmup import zip_inputs
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    LaunchSpec,
+    TritonWarmupTensor,
+    VllmTritonJitKernel,
+    kernel_launcher,
+    triton_scalar_specialization_rep,
+)
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
 
@@ -17,49 +28,124 @@ def get_dspark_swa_index_width(
     return cdiv(width, _DSPARK_SWA_INDEX_ALIGNMENT) * _DSPARK_SWA_INDEX_ALIGNMENT
 
 
-@triton.jit
-def _compressed_slot_mapping_kernel(
-    # [num_tokens]
-    slot_mapping_ptr,
-    # [num_reqs + 1]
-    query_start_loc_ptr,
-    # [num_reqs]
-    seq_lens_ptr,
-    # [num_reqs, max_num_blocks]
-    block_table_ptr,
-    block_table_stride,
-    block_size,
-    COMPRESS_RATIO: tl.constexpr,
-    PAD_ID: tl.constexpr,
-    TRITON_BLOCK_SIZE: tl.constexpr,
+class CompressedSlotMappingKernel(
+    VllmTritonJitKernel["CompressedSlotMappingKernel.CompileKey"]
 ):
-    batch_idx = tl.program_id(0)
+    TRITON_BLOCK_SIZE = 1024
 
-    query_start = tl.load(query_start_loc_ptr + batch_idx)
-    query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
-    query_len = query_end - query_start
+    @dataclass(frozen=True)
+    class CompileKey:
+        compress_ratio: int
+        triton_block_size: int
+        block_size: int
 
-    seq_len = tl.load(seq_lens_ptr + batch_idx)
-    start_pos = seq_len - query_len
+    @staticmethod
+    @triton.jit(do_not_specialize=["block_table_stride"])
+    def kernel(
+        # [num_tokens]
+        slot_mapping_ptr,
+        # [num_reqs + 1]
+        query_start_loc_ptr,
+        # [num_reqs]
+        seq_lens_ptr,
+        # [num_reqs, max_num_blocks]
+        block_table_ptr,
+        block_table_stride,
+        block_size,
+        COMPRESS_RATIO: tl.constexpr,
+        PAD_ID: tl.constexpr,
+        TRITON_BLOCK_SIZE: tl.constexpr,
+    ):
+        batch_idx = tl.program_id(0)
 
-    for i in range(0, query_len, TRITON_BLOCK_SIZE):
-        offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
-        mask = offset < query_len
+        query_start = tl.load(query_start_loc_ptr + batch_idx)
+        query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
+        query_len = query_end - query_start
 
-        pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)
-        is_valid = (pos + 1) % COMPRESS_RATIO == 0
-        pos_after_compress = pos // COMPRESS_RATIO
+        seq_len = tl.load(seq_lens_ptr + batch_idx)
+        start_pos = seq_len - query_len
 
-        block_ids = pos_after_compress // block_size
-        block_numbers = tl.load(
-            block_table_ptr + batch_idx * block_table_stride + block_ids,
-            mask=mask & is_valid,
+        for i in range(0, query_len, TRITON_BLOCK_SIZE):
+            offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+            mask = offset < query_len
+
+            pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE)
+            is_valid = (pos + 1) % COMPRESS_RATIO == 0
+            pos_after_compress = pos // COMPRESS_RATIO
+
+            block_ids = pos_after_compress // block_size
+            block_numbers = tl.load(
+                block_table_ptr + batch_idx * block_table_stride + block_ids,
+                mask=mask & is_valid,
+            )
+            slot_ids = block_numbers * block_size + pos_after_compress % block_size
+
+            # NOTE
+            slot_ids = tl.where(is_valid, slot_ids, PAD_ID)
+            tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        compress_ratio: int,
+        block_size: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            compress_ratio=compress_ratio,
+            triton_block_size=self.TRITON_BLOCK_SIZE,
+            block_size=triton_scalar_specialization_rep(block_size),
         )
-        slot_ids = block_numbers * block_size + pos_after_compress % block_size
 
-        # NOTE
-        slot_ids = tl.where(is_valid, slot_ids, PAD_ID)
-        tl.store(slot_mapping_ptr + query_start + offset, slot_ids, mask=mask)
+    def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]:
+        hf_config = vllm_config.model_config.hf_config
+        configured_ratios = (
+            *(getattr(hf_config, "compress_ratios", None) or ()),
+            getattr(hf_config, "index_kpool", 1) or 1,
+        )
+        compress_ratios = tuple(
+            dict.fromkeys(int(ratio) for ratio in configured_ratios if int(ratio) > 1)
+        )
+        if not compress_ratios:
+            return []
+        return self._trace_dispatch(self.dispatch)(
+            zip_inputs(
+                *(
+                    dict(
+                        compress_ratio=ratio,
+                        block_size=vllm_config.cache_config.block_size // ratio,
+                    )
+                    for ratio in compress_ratios
+                )
+            )
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        return dict(
+            slot_mapping=TritonWarmupTensor(torch.int64),
+            query_start_loc=int32_ptr,
+            seq_lens=int32_ptr,
+            block_table=int32_ptr,
+            block_size=compile_key.block_size,
+            compress_ratio=compile_key.compress_ratio,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        slot_mapping: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        seq_lens: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        compress_ratio: int,
+    ) -> LaunchSpec:
+        return (block_table.shape[0],), dict(
+            block_table_stride=block_table.stride(0),
+            COMPRESS_RATIO=compress_ratio,
+            PAD_ID=-1,
+            TRITON_BLOCK_SIZE=self.TRITON_BLOCK_SIZE,
+        )
 
 
 def get_compressed_slot_mapping(
@@ -83,16 +169,15 @@ def get_compressed_slot_mapping(
             (num_tokens,), -1, dtype=torch.int64, device=query_start_loc.device
         )
 
-    num_reqs = block_table.shape[0]
-    _compressed_slot_mapping_kernel[(num_reqs,)](
+    _COMPRESSED_SLOT_MAPPING_KERNEL(
         slot_mapping,
         query_start_loc,
         seq_lens,
         block_table,
-        block_table.stride(0),
         block_size,
         compress_ratio,
-        PAD_ID=-1,
-        TRITON_BLOCK_SIZE=1024,
     )
     return slot_mapping
+
+
+_COMPRESSED_SLOT_MAPPING_KERNEL = CompressedSlotMappingKernel()
diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py
index 4700fbea06..206023f5da 100644
--- a/vllm/v1/attention/backends/mla/indexer.py
+++ b/vllm/v1/attention/backends/mla/indexer.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
+from typing import Any
 
 import torch
 
@@ -8,13 +9,13 @@ import vllm.envs as envs
 from vllm.config import VllmConfig
 from vllm.distributed import get_dcp_group, get_pcp_group
 from vllm.logger import init_logger
-from vllm.model_executor.warmup.jit_warmup import (
-    VllmJitKernel,
-    WarmupIntRange,
-)
 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    LaunchSpec,
     TritonPointerInputVariant,
     TritonWarmupTensor,
+    VllmTritonJitKernel,
+    kernel_launcher,
+    triton_scalar_specialization_rep,
 )
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
@@ -63,41 +64,112 @@ def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool:
     return use_fp4
 
 
-@triton.jit
-def _prepare_uniform_decode_kernel(
-    seq_lens_ptr,
-    decode_seq_lens_ptr,
-    block_table_ptr,
-    block_table_stride,
-    expanded_block_table_ptr,
-    expanded_bt_stride,
-    decode_lens_ptr,
-    max_decode_len,
-    BLOCK_SIZE: tl.constexpr,
+class PrepareUniformDecodeKernel(
+    VllmTritonJitKernel["PrepareUniformDecodeKernel.CompileKey"]
 ):
-    idx = tl.program_id(0)
-    req_id = idx // max_decode_len
-    local_idx = idx % max_decode_len
+    BLOCK_SIZE = 1024
 
-    # Compute number of KVs attended to by this token. Padding requests have
-    # seq_len == 0, which would otherwise make the first token of each padded
-    # request negative (e.g. next_n=2 gives 0-2+0+1 = -1). Downstream kernels
-    # read these as uint32, turning -1 into ~4e9.
-    seq_len = tl.load(seq_lens_ptr + req_id)
-    per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0)
-    tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
+    @dataclass(frozen=True)
+    class CompileKey:
+        block_size: int
+        block_table_stride: int
+        expanded_bt_stride: int
+        max_decode_len: int
 
-    # Copy block table row.
-    src = block_table_ptr + req_id * block_table_stride
-    dst = expanded_block_table_ptr + idx * expanded_bt_stride
-    for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
-        off = i + tl.arange(0, BLOCK_SIZE)
-        mask = off < expanded_bt_stride
-        src_block = tl.load(src + off, mask=mask)
-        tl.store(dst + off, src_block, mask=mask)
+    @staticmethod
+    @triton.jit
+    def kernel(
+        seq_lens_ptr,
+        decode_seq_lens_ptr,
+        block_table_ptr,
+        block_table_stride,
+        expanded_block_table_ptr,
+        expanded_bt_stride,
+        decode_lens_ptr,
+        max_decode_len,
+        BLOCK_SIZE: tl.constexpr,
+    ):
+        idx = tl.program_id(0)
+        req_id = idx // max_decode_len
+        local_idx = idx % max_decode_len
 
-    # All reqs now have decode_len = 1.
-    tl.store(decode_lens_ptr + idx, 1)
+        # Compute number of KVs attended to by this token. Padding requests have
+        # seq_len == 0, which would otherwise make the first token of each padded
+        # request negative (e.g. next_n=2 gives 0-2+0+1 = -1). Downstream kernels
+        # read these as uint32, turning -1 into ~4e9.
+        seq_len = tl.load(seq_lens_ptr + req_id)
+        per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0)
+        tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
+
+        # Copy block table row.
+        src = block_table_ptr + req_id * block_table_stride
+        dst = expanded_block_table_ptr + idx * expanded_bt_stride
+        for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
+            off = i + tl.arange(0, BLOCK_SIZE)
+            mask = off < expanded_bt_stride
+            src_block = tl.load(src + off, mask=mask)
+            tl.store(dst + off, src_block, mask=mask)
+
+        # All reqs now have decode_len = 1.
+        tl.store(decode_lens_ptr + idx, 1)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        block_size: int,
+        block_table_stride: int,
+        expanded_bt_stride: int,
+        max_decode_len: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            block_size=block_size,
+            block_table_stride=triton_scalar_specialization_rep(block_table_stride),
+            expanded_bt_stride=triton_scalar_specialization_rep(expanded_bt_stride),
+            max_decode_len=triton_scalar_specialization_rep(max_decode_len),
+        )
+
+    def get_warmup_keys(self) -> list[CompileKey]:
+        return self._trace_dispatch(self.dispatch)(
+            block_size=self.BLOCK_SIZE,
+            block_table_stride=(1, 2, 16),
+            expanded_bt_stride=(1, 2, 16),
+            max_decode_len=(1, 2, 16),
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        return dict(
+            seq_lens=int32_ptr,
+            decode_seq_lens=int32_ptr,
+            block_table=TritonWarmupTensor(
+                torch.int32,
+                shape=(1, compile_key.block_table_stride),
+            ),
+            expanded_block_table=TritonWarmupTensor(
+                torch.int32,
+                shape=(1, compile_key.expanded_bt_stride),
+            ),
+            decode_lens=int32_ptr,
+            num_decode_tokens=1,
+            max_decode_len=compile_key.max_decode_len,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        seq_lens: torch.Tensor,
+        decode_seq_lens: torch.Tensor,
+        block_table: torch.Tensor,
+        expanded_block_table: torch.Tensor,
+        decode_lens: torch.Tensor,
+        num_decode_tokens: int,
+        max_decode_len: int,
+    ) -> LaunchSpec:
+        return (num_decode_tokens,), dict(
+            block_table_stride=block_table.stride(0),
+            expanded_bt_stride=expanded_block_table.stride(0),
+            BLOCK_SIZE=self.BLOCK_SIZE,
+        )
 
 
 def split_indexer_prefill_chunks(
@@ -215,14 +287,8 @@ class DeepseekV32IndexerPrefillChunkMetadata:
     max_local_total_seq_lens: int = 0
 
 
-_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS = (
-    TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),
-    TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),
-)
-
-
 class BuildPrefillChunkMetadataKernel(
-    VllmJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]
+    VllmTritonJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"]
 ):
     BLOCK_SIZE = 1024
 
@@ -230,11 +296,11 @@ class BuildPrefillChunkMetadataKernel(
     class CompileKey:
         query_slice_start: int
         query_slice_stop: int
-        DCP_RANK: int
-        DCP_WORLD: int
-        DCP_INTERLEAVE: int
-        BLOCK_SIZE: int
-        COMPRESS_RATIO: int
+        dcp_rank: int
+        dcp_world: int
+        dcp_interleave: int
+        block_size: int
+        compress_ratio: int
         input_variant: TritonPointerInputVariant
 
     @staticmethod
@@ -327,13 +393,13 @@ class BuildPrefillChunkMetadataKernel(
         input_variant: TritonPointerInputVariant,
     ) -> CompileKey:
         return self.CompileKey(
-            query_slice_start=query_slice_start,
-            query_slice_stop=query_slice_stop,
-            DCP_RANK=DCP_RANK,
-            DCP_WORLD=DCP_WORLD,
-            DCP_INTERLEAVE=DCP_INTERLEAVE,
-            BLOCK_SIZE=BLOCK_SIZE,
-            COMPRESS_RATIO=COMPRESS_RATIO,
+            query_slice_start=triton_scalar_specialization_rep(query_slice_start),
+            query_slice_stop=triton_scalar_specialization_rep(query_slice_stop),
+            dcp_rank=triton_scalar_specialization_rep(DCP_RANK),
+            dcp_world=triton_scalar_specialization_rep(DCP_WORLD),
+            dcp_interleave=triton_scalar_specialization_rep(DCP_INTERLEAVE),
+            block_size=BLOCK_SIZE,
+            compress_ratio=COMPRESS_RATIO,
             input_variant=input_variant,
         )
 
@@ -345,42 +411,51 @@ class BuildPrefillChunkMetadataKernel(
         dcp_interleave = parallel_config.cp_kv_cache_interleave_size
         dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0
         compress_ratios = tuple(
-            max(1, int(ratio))
-            for ratio in (getattr(hf_config, "compress_ratios", None) or (1,))
+            dict.fromkeys(
+                max(1, int(ratio))
+                for ratio in (
+                    *(getattr(hf_config, "compress_ratios", None) or (1,)),
+                    getattr(hf_config, "index_kpool", 1) or 1,
+                )
+            )
         )
         return self._trace_dispatch(self.dispatch)(
-            query_slice_start=WarmupIntRange(0, 2),
+            # Cover Triton's divisible, exact-one, and generic i32 classes.
+            query_slice_start=(0, 1, 2),
             query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens),
             DCP_RANK=dcp_rank,
             DCP_WORLD=dcp_world,
             DCP_INTERLEAVE=dcp_interleave,
             BLOCK_SIZE=self.BLOCK_SIZE,
             COMPRESS_RATIO=list(compress_ratios),
-            input_variant=_BUILD_PREFILL_CHUNK_METADATA_INPUT_VARIANTS,
+            input_variant=(
+                TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=True),
+                TritonPointerInputVariant.from_alignment(uncompressed_seq_lens=False),
+            ),
         )
 
-    def compile(self, compile_key: CompileKey) -> None:
-        warmup = getattr(self.kernel, "warmup", None)
-        assert warmup is not None
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
         int32_ptr = TritonWarmupTensor(torch.int32)
-        warmup(
-            int32_ptr,
-            compile_key.input_variant.pointer("uncompressed_seq_lens", torch.int32),
-            int32_ptr,
-            int32_ptr,
-            int32_ptr,
-            int32_ptr,
-            int32_ptr,
-            compile_key.query_slice_start,
-            compile_key.query_slice_stop,
-            compile_key.DCP_RANK,
-            compile_key.DCP_WORLD,
-            compile_key.DCP_INTERLEAVE,
-            BLOCK_SIZE=compile_key.BLOCK_SIZE,
-            COMPRESS_RATIO=compile_key.COMPRESS_RATIO,
-            grid=(1,),
+        return dict(
+            query_start_loc=int32_ptr,
+            uncompressed_seq_lens=compile_key.input_variant.pointer(
+                "uncompressed_seq_lens", torch.int32
+            ),
+            cu_compressed_seq_lens=int32_ptr,
+            row_start_cu_compressed_seq_lens=int32_ptr,
+            token_to_seq=int32_ptr,
+            cu_compressed_seq_len_ks=int32_ptr,
+            cu_compressed_seq_len_ke=int32_ptr,
+            query_slice_start=compile_key.query_slice_start,
+            query_slice_stop=compile_key.query_slice_stop,
+            DCP_RANK=compile_key.dcp_rank,
+            DCP_WORLD=compile_key.dcp_world,
+            DCP_INTERLEAVE=compile_key.dcp_interleave,
+            num_reqs=1,
+            COMPRESS_RATIO=compile_key.compress_ratio,
         )
 
+    @kernel_launcher
     def __call__(
         self,
         query_start_loc: torch.Tensor,
@@ -398,28 +473,13 @@ class BuildPrefillChunkMetadataKernel(
         *,
         num_reqs: int,
         COMPRESS_RATIO: int,
-    ) -> None:
-        self.kernel[(num_reqs,)](
-            query_start_loc,
-            uncompressed_seq_lens,
-            cu_compressed_seq_lens,
-            row_start_cu_compressed_seq_lens,
-            token_to_seq,
-            cu_compressed_seq_len_ks,
-            cu_compressed_seq_len_ke,
-            query_slice_start,
-            query_slice_stop,
-            DCP_RANK,
-            DCP_WORLD,
-            DCP_INTERLEAVE,
+    ) -> LaunchSpec:
+        return (num_reqs,), dict(
             BLOCK_SIZE=self.BLOCK_SIZE,
             COMPRESS_RATIO=COMPRESS_RATIO,
         )
 
 
-_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()
-
-
 @dataclass
 class DeepseekV32IndexerPrefillMetadata:
     chunks: list[DeepseekV32IndexerPrefillChunkMetadata]
@@ -694,16 +754,14 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
                 and num_decodes * max_decode_len == num_decode_tokens
             ):
                 # Uniform decode lengths with no cudagraph token padding.
-                _prepare_uniform_decode_kernel[(num_decode_tokens,)](
+                _PREPARE_UNIFORM_DECODE_KERNEL(
                     seq_lens,
                     self.decode_seq_lens_buffer,
                     block_table,
-                    block_table.stride(0),
                     self.expanded_block_table_buffer,
-                    self.expanded_block_table_buffer.stride(0),
                     self.decode_lens_buffer,
+                    num_decode_tokens,
                     max_decode_len,
-                    BLOCK_SIZE=1024,
                 )
                 self.decode_seq_lens_buffer[num_decode_tokens:] = 0
                 seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
@@ -1180,3 +1238,7 @@ def build_prefill_chunk_metadata(
         local_total_seq_lens=local_total_seq_lens,
         max_local_total_seq_lens=max_local_total_seq_lens,
     )
+
+
+_PREPARE_UNIFORM_DECODE_KERNEL = PrepareUniformDecodeKernel()
+_BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel()
diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py
index 83b1afef1f..b968bda9fd 100644
--- a/vllm/v1/attention/backends/mla/sparse_swa.py
+++ b/vllm/v1/attention/backends/mla/sparse_swa.py
@@ -1,17 +1,21 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass, field
-from typing import ClassVar, cast
+from typing import Any, ClassVar, cast
 
 import torch
 
 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.warmup.jit_warmup import (
-    VllmJitKernel,
     WarmupIntRange,
 )
-from vllm.model_executor.warmup.jit_warmup_triton_helper import TritonWarmupTensor
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    LaunchSpec,
+    TritonWarmupTensor,
+    VllmTritonJitKernel,
+    kernel_launcher,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv, next_power_of_2
@@ -275,11 +279,11 @@ class DeepseekSparseSWAMetadata:
 
 
 class ComputePrefillMetadataKernel(
-    VllmJitKernel["ComputePrefillMetadataKernel.CompileKey"]
+    VllmTritonJitKernel["ComputePrefillMetadataKernel.CompileKey"]
 ):
     @dataclass(frozen=True)
     class CompileKey:
-        BLOCK_SIZE: int
+        block_size: int
 
     @staticmethod
     @triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"])
@@ -320,7 +324,7 @@ class ComputePrefillMetadataKernel(
         num_prefills: int,
     ) -> CompileKey:
         return self.CompileKey(
-            BLOCK_SIZE=next_power_of_2(num_prefills),
+            block_size=next_power_of_2(num_prefills),
         )
 
     def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]:
@@ -336,21 +340,18 @@ class ComputePrefillMetadataKernel(
             num_prefills=WarmupIntRange(1, max_prefills + 1),
         )
 
-    def compile(self, compile_key: CompileKey) -> None:
-        warmup = getattr(self.kernel, "warmup", None)
-        assert warmup is not None
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
         int32_ptr = TritonWarmupTensor(torch.int32)
-        warmup(
-            int32_ptr,
-            int32_ptr,
-            int32_ptr,
-            compile_key.BLOCK_SIZE,
-            0,
-            1,
-            BLOCK_SIZE=compile_key.BLOCK_SIZE,
-            grid=(1,),
+        return dict(
+            prefill_gather_lens=int32_ptr,
+            seq_lens=int32_ptr,
+            query_start_loc=int32_ptr,
+            num_prefills=compile_key.block_size,
+            num_decodes=0,
+            window_size=1,
         )
 
+    @kernel_launcher
     def __call__(
         self,
         prefill_gather_lens: torch.Tensor,
@@ -359,22 +360,13 @@ class ComputePrefillMetadataKernel(
         num_prefills: int,
         num_decodes: int,
         window_size: int,
-    ) -> None:
+    ) -> LaunchSpec:
         compile_key = self.dispatch(num_prefills=num_prefills)
-        self.kernel[(1,)](
-            prefill_gather_lens,
-            seq_lens,
-            query_start_loc,
-            num_prefills,
-            num_decodes,
-            window_size,
-            BLOCK_SIZE=compile_key.BLOCK_SIZE,
+        return (1,), dict(
+            BLOCK_SIZE=compile_key.block_size,
         )
 
 
-_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel()
-
-
 class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
     """Builds metadata for DeepseekV4 SWA cache.
 
@@ -545,9 +537,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
                         device=self.device,
                     )
                 decode_swa_indices = self.decode_swa_indices_noncausal
-                _compute_dspark_noncausal_swa_indices_kernel[(num_decode_tokens,)](
+                _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL(
                     decode_swa_indices,
-                    decode_swa_indices.stride(0),
                     self.decode_swa_lens,
                     self.window_size,
                     self.noncausal_index_width,
@@ -556,15 +547,13 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
                     token_to_req_indices,
                     is_valid_token,
                     block_table,
-                    block_table.stride(0),
                     self.block_size,
+                    num_tokens=num_decode_tokens,
                     token_offset=0,
-                    TRITON_BLOCK_SIZE=1024,
                 )
             else:
-                _compute_swa_indices_and_lens_kernel[(num_decode_tokens,)](
+                _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
                     decode_swa_indices,
-                    decode_swa_indices.stride(0),
                     self.decode_swa_lens,
                     self.window_size,
                     query_start_loc,
@@ -572,10 +561,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
                     token_to_req_indices,
                     is_valid_token,
                     block_table,
-                    block_table.stride(0),
                     self.block_size,
+                    num_tokens=num_decode_tokens,
                     token_offset=0,
-                    TRITON_BLOCK_SIZE=1024,
                 )
 
         # Prefill SWA indices live in paged coordinates. `token_offset` lets
@@ -584,9 +572,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
         if num_prefill_tokens > 0:
             prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens]
             prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens]
-            _compute_swa_indices_and_lens_kernel[(num_prefill_tokens,)](
+            _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
                 prefill_swa_indices,
-                prefill_swa_indices.stride(0),
                 prefill_swa_lens,
                 self.window_size,
                 query_start_loc,
@@ -594,10 +581,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
                 token_to_req_indices,
                 is_valid_token,
                 block_table,
-                block_table.stride(0),
                 self.block_size,
+                num_tokens=num_prefill_tokens,
                 token_offset=num_decode_tokens,
-                TRITON_BLOCK_SIZE=1024,
             )
 
         # Pre-compute DeepseekV4 prefill metadata shared across all attention layers.
@@ -668,9 +654,8 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
         assert metadata.decode_swa_indices is not None
         assert metadata.decode_swa_lens is not None
 
-        _compute_swa_indices_and_lens_kernel[(metadata.num_decode_tokens,)](
+        _COMPUTE_SWA_INDICES_AND_LENS_KERNEL(
             metadata.decode_swa_indices,
-            metadata.decode_swa_indices.stride(0),
             metadata.decode_swa_lens,
             metadata.decode_swa_indices.shape[-1],
             metadata.query_start_loc,
@@ -678,10 +663,9 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
             metadata.token_to_req_indices,
             metadata.is_valid_token,
             metadata.block_table,
-            metadata.block_table.stride(0),
             self.block_size,
+            num_tokens=metadata.num_decode_tokens,
             token_offset=0,
-            TRITON_BLOCK_SIZE=1024,
         )
         tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens)
         metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY]
@@ -771,142 +755,307 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
         return result
 
 
-@triton.jit(do_not_specialize=["token_offset"])
-def _compute_swa_indices_and_lens_kernel(
-    swa_indices_ptr,
-    swa_indices_stride,
-    swa_lens_ptr,
-    window_size,
-    query_start_loc_ptr,
-    seq_lens_ptr,
-    token_to_req_indices_ptr,
-    is_valid_token_ptr,
-    block_table_ptr,
-    block_table_stride,
-    block_size,
-    token_offset,
-    TRITON_BLOCK_SIZE: tl.constexpr,
+# TODO(ben): unify this kernel to reduce duplication
+
+
+class ComputeSWAIndicesAndLensKernel(
+    VllmTritonJitKernel["ComputeSWAIndicesAndLensKernel.CompileKey"]
 ):
-    pid = tl.program_id(0)
-    token_idx = pid + token_offset
-    is_valid = tl.load(is_valid_token_ptr + token_idx)
-    if not is_valid:
-        tl.store(swa_lens_ptr + pid, 0)
-        # Clear the row so a padded token cannot gather through stale indices.
+    @dataclass(frozen=True)
+    class CompileKey:
+        window_size: int
+        block_size: int
+        triton_block_size: int
+
+    @staticmethod
+    @triton.jit(
+        do_not_specialize=[
+            "swa_indices_stride",
+            "block_table_stride",
+            "token_offset",
+        ]
+    )
+    def kernel(
+        swa_indices_ptr,
+        swa_indices_stride,
+        swa_lens_ptr,
+        window_size,
+        query_start_loc_ptr,
+        seq_lens_ptr,
+        token_to_req_indices_ptr,
+        is_valid_token_ptr,
+        block_table_ptr,
+        block_table_stride,
+        block_size,
+        token_offset,
+        TRITON_BLOCK_SIZE: tl.constexpr,
+    ):
+        pid = tl.program_id(0)
+        token_idx = pid + token_offset
+        is_valid = tl.load(is_valid_token_ptr + token_idx)
+        if not is_valid:
+            tl.store(swa_lens_ptr + pid, 0)
+            # Clear the row so a padded token cannot gather through stale indices.
+            for i in range(0, window_size, TRITON_BLOCK_SIZE):
+                offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+                tl.store(
+                    swa_indices_ptr + pid * swa_indices_stride + offset,
+                    -1,
+                    mask=offset < window_size,
+                )
+            return
+
+        req_idx = tl.load(token_to_req_indices_ptr + token_idx)
+
+        query_start = tl.load(query_start_loc_ptr + req_idx)
+        query_end = tl.load(query_start_loc_ptr + req_idx + 1)
+        query_len = query_end - query_start
+
+        seq_len = tl.load(seq_lens_ptr + req_idx)
+        prefix_len = seq_len - query_len
+
+        pos = prefix_len + token_idx - query_start
+        start_pos = tl.maximum(pos - window_size + 1, 0)
+        end_pos = pos + 1
+
+        swa_len = end_pos - start_pos
+        tl.store(swa_lens_ptr + pid, swa_len)
+
         for i in range(0, window_size, TRITON_BLOCK_SIZE):
             offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+
+            pos_offset = start_pos + offset
+            block_indices = pos_offset // block_size
+            block_numbers = tl.load(
+                block_table_ptr + req_idx * block_table_stride + block_indices,
+                mask=pos_offset < end_pos,
+            )
+            block_offsets = pos_offset % block_size
+            slot_ids = block_numbers * block_size + block_offsets
+
+            slot_ids = tl.where(offset < swa_len, slot_ids, -1)
             tl.store(
                 swa_indices_ptr + pid * swa_indices_stride + offset,
-                -1,
+                slot_ids,
                 mask=offset < window_size,
             )
-        return
 
-    req_idx = tl.load(token_to_req_indices_ptr + token_idx)
-
-    query_start = tl.load(query_start_loc_ptr + req_idx)
-    query_end = tl.load(query_start_loc_ptr + req_idx + 1)
-    query_len = query_end - query_start
-
-    seq_len = tl.load(seq_lens_ptr + req_idx)
-    prefix_len = seq_len - query_len
-
-    pos = prefix_len + token_idx - query_start
-    start_pos = tl.maximum(pos - window_size + 1, 0)
-    end_pos = pos + 1
-
-    swa_len = end_pos - start_pos
-    tl.store(swa_lens_ptr + pid, swa_len)
-
-    for i in range(0, window_size, TRITON_BLOCK_SIZE):
-        offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
-
-        pos_offset = start_pos + offset
-        block_indices = pos_offset // block_size
-        block_numbers = tl.load(
-            block_table_ptr + req_idx * block_table_stride + block_indices,
-            mask=pos_offset < end_pos,
+    def dispatch(  # type: ignore[override]
+        self,
+        **compile_key_fields: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            **compile_key_fields,
+            triton_block_size=1024,
         )
-        block_offsets = pos_offset % block_size
-        slot_ids = block_numbers * block_size + block_offsets
 
-        slot_ids = tl.where(offset < swa_len, slot_ids, -1)
-        tl.store(
-            swa_indices_ptr + pid * swa_indices_stride + offset,
-            slot_ids,
-            mask=offset < window_size,
+    def get_warmup_keys(
+        self,
+        *,
+        window_size: int,
+        block_size: int,
+    ) -> list[CompileKey]:
+        return self._trace_dispatch(self.dispatch)(
+            window_size=window_size,
+            block_size=block_size,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        return dict(
+            swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
+            swa_lens=int32_ptr,
+            window_size=compile_key.window_size,
+            query_start_loc=int32_ptr,
+            seq_lens=int32_ptr,
+            token_to_req_indices=int32_ptr,
+            is_valid_token=TritonWarmupTensor(torch.bool),
+            block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
+            block_size=compile_key.block_size,
+            num_tokens=1,
+            token_offset=0,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        swa_indices: torch.Tensor,
+        swa_lens: torch.Tensor,
+        window_size: int,
+        query_start_loc: torch.Tensor,
+        seq_lens: torch.Tensor,
+        token_to_req_indices: torch.Tensor,
+        is_valid_token: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        *,
+        num_tokens: int,
+        token_offset: int,
+    ) -> LaunchSpec:
+        return (num_tokens,), dict(
+            swa_indices_stride=swa_indices.stride(0),
+            block_table_stride=block_table.stride(0),
+            TRITON_BLOCK_SIZE=1024,
         )
 
 
-# TODO(ben): unify this kernel to reduce duplication
-@triton.jit(do_not_specialize=["token_offset"])
-def _compute_dspark_noncausal_swa_indices_kernel(
-    swa_indices_ptr,
-    swa_indices_stride,
-    swa_lens_ptr,
-    window_size,
-    index_width,
-    query_start_loc_ptr,
-    seq_lens_ptr,
-    token_to_req_indices_ptr,
-    is_valid_token_ptr,
-    block_table_ptr,
-    block_table_stride,
-    block_size,
-    token_offset,
-    TRITON_BLOCK_SIZE: tl.constexpr,
+class ComputeDSparkNoncausalSWAIndicesKernel(
+    VllmTritonJitKernel["ComputeDSparkNoncausalSWAIndicesKernel.CompileKey"]
 ):
-    """Non-causal per-token indices for the DSpark draft block.
+    @dataclass(frozen=True)
+    class CompileKey:
+        window_size: int
+        index_width: int
+        block_size: int
+        triton_block_size: int
+
+    @staticmethod
+    @triton.jit(
+        do_not_specialize=[
+            "swa_indices_stride",
+            "block_table_stride",
+            "token_offset",
+        ]
+    )
+    def kernel(
+        swa_indices_ptr,
+        swa_indices_stride,
+        swa_lens_ptr,
+        window_size,
+        index_width,
+        query_start_loc_ptr,
+        seq_lens_ptr,
+        token_to_req_indices_ptr,
+        is_valid_token_ptr,
+        block_table_ptr,
+        block_table_stride,
+        block_size,
+        token_offset,
+        TRITON_BLOCK_SIZE: tl.constexpr,
+    ):
+        """Non-causal per-token indices for the DSpark draft block.
+
+        Here, we populate the topk indices with the trailing window of context tokens,
+        plus all query tokens (including future ones).
+        """
+        pid = tl.program_id(0)
+        token_idx = pid + token_offset
+        is_valid = tl.load(is_valid_token_ptr + token_idx)
+        if not is_valid:
+            tl.store(swa_lens_ptr + pid, 0)
+            # Clear the row so a padded token cannot gather through stale indices.
+            for i in range(0, index_width, TRITON_BLOCK_SIZE):
+                offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+                tl.store(
+                    swa_indices_ptr + pid * swa_indices_stride + offset,
+                    -1,
+                    mask=offset < index_width,
+                )
+            return
+
+        req_idx = tl.load(token_to_req_indices_ptr + token_idx)
+
+        query_start = tl.load(query_start_loc_ptr + req_idx)
+        query_end = tl.load(query_start_loc_ptr + req_idx + 1)
+        query_len = query_end - query_start
+
+        seq_len = tl.load(seq_lens_ptr + req_idx)
+        prefix_len = seq_len - query_len
+
+        # Block-anchored window (shared by every token in the block) + full block.
+        start_pos = tl.maximum(prefix_len - window_size, 0)
+        end_pos = seq_len
+
+        swa_len = end_pos - start_pos
+        tl.store(swa_lens_ptr + pid, swa_len)
 
-    Here, we populate the topk indices with the trailing window of context tokens,
-    plus all query tokens (including future ones).
-    """
-    pid = tl.program_id(0)
-    token_idx = pid + token_offset
-    is_valid = tl.load(is_valid_token_ptr + token_idx)
-    if not is_valid:
-        tl.store(swa_lens_ptr + pid, 0)
-        # Clear the row so a padded token cannot gather through stale indices.
         for i in range(0, index_width, TRITON_BLOCK_SIZE):
             offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
+
+            pos_offset = start_pos + offset
+            block_indices = pos_offset // block_size
+            block_numbers = tl.load(
+                block_table_ptr + req_idx * block_table_stride + block_indices,
+                mask=pos_offset < end_pos,
+            )
+            block_offsets = pos_offset % block_size
+            slot_ids = block_numbers * block_size + block_offsets
+
+            slot_ids = tl.where(offset < swa_len, slot_ids, -1)
             tl.store(
                 swa_indices_ptr + pid * swa_indices_stride + offset,
-                -1,
+                slot_ids,
                 mask=offset < index_width,
             )
-        return
 
-    req_idx = tl.load(token_to_req_indices_ptr + token_idx)
-
-    query_start = tl.load(query_start_loc_ptr + req_idx)
-    query_end = tl.load(query_start_loc_ptr + req_idx + 1)
-    query_len = query_end - query_start
-
-    seq_len = tl.load(seq_lens_ptr + req_idx)
-    prefix_len = seq_len - query_len
-
-    # Block-anchored window (shared by every token in the block) + full block.
-    start_pos = tl.maximum(prefix_len - window_size, 0)
-    end_pos = seq_len
-
-    swa_len = end_pos - start_pos
-    tl.store(swa_lens_ptr + pid, swa_len)
-
-    for i in range(0, index_width, TRITON_BLOCK_SIZE):
-        offset = i + tl.arange(0, TRITON_BLOCK_SIZE)
-
-        pos_offset = start_pos + offset
-        block_indices = pos_offset // block_size
-        block_numbers = tl.load(
-            block_table_ptr + req_idx * block_table_stride + block_indices,
-            mask=pos_offset < end_pos,
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        window_size: int,
+        num_speculative_tokens: int,
+        block_size: int,
+    ) -> CompileKey:
+        return self.CompileKey(
+            window_size=window_size,
+            index_width=get_dspark_swa_index_width(window_size, num_speculative_tokens),
+            block_size=block_size,
+            triton_block_size=1024,
         )
-        block_offsets = pos_offset % block_size
-        slot_ids = block_numbers * block_size + block_offsets
 
-        slot_ids = tl.where(offset < swa_len, slot_ids, -1)
-        tl.store(
-            swa_indices_ptr + pid * swa_indices_stride + offset,
-            slot_ids,
-            mask=offset < index_width,
+    def get_warmup_keys(
+        self,
+        *,
+        window_size: int,
+        num_speculative_tokens: int,
+        block_size: int,
+    ) -> list[CompileKey]:
+        return self._trace_dispatch(self.dispatch)(
+            window_size=window_size,
+            num_speculative_tokens=num_speculative_tokens,
+            block_size=block_size,
         )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        return dict(
+            swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
+            swa_lens=int32_ptr,
+            window_size=compile_key.window_size,
+            index_width=compile_key.index_width,
+            query_start_loc=int32_ptr,
+            seq_lens=int32_ptr,
+            token_to_req_indices=int32_ptr,
+            is_valid_token=TritonWarmupTensor(torch.bool),
+            block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)),
+            block_size=compile_key.block_size,
+            num_tokens=1,
+            token_offset=0,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        swa_indices: torch.Tensor,
+        swa_lens: torch.Tensor,
+        window_size: int,
+        index_width: int,
+        query_start_loc: torch.Tensor,
+        seq_lens: torch.Tensor,
+        token_to_req_indices: torch.Tensor,
+        is_valid_token: torch.Tensor,
+        block_table: torch.Tensor,
+        block_size: int,
+        *,
+        num_tokens: int,
+        token_offset: int,
+    ) -> LaunchSpec:
+        return (num_tokens,), dict(
+            swa_indices_stride=swa_indices.stride(0),
+            block_table_stride=block_table.stride(0),
+            TRITON_BLOCK_SIZE=1024,
+        )
+
+
+_COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel()
+_COMPUTE_SWA_INDICES_AND_LENS_KERNEL = ComputeSWAIndicesAndLensKernel()
+_COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL = ComputeDSparkNoncausalSWAIndicesKernel()
diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py
index 00679d5e26..168a8c05c1 100644
--- a/vllm/v1/attention/backends/mla/sparse_utils.py
+++ b/vllm/v1/attention/backends/mla/sparse_utils.py
@@ -2,9 +2,26 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Utility functions for sparse MLA backends."""
 
+from dataclasses import dataclass
+from typing import Any
+
 import torch
 
+from vllm.config import VllmConfig
+from vllm.distributed import get_dcp_group
+from vllm.model_executor.warmup.jit_warmup import (
+    zip_inputs,
+)
+from vllm.model_executor.warmup.jit_warmup_triton_helper import (
+    LaunchSpec,
+    TritonWarmupTensor,
+    VllmTritonJitKernel,
+    kernel_launcher,
+    triton_scalar_specialization_rep,
+)
 from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import cdiv
+from vllm.v1.worker.block_table import get_block_table_width
 
 
 def flat_kv_row_view(
@@ -30,126 +47,330 @@ def flat_kv_row_view(
 
 
 # Kernel with prefill workspace support and valid count tracking
-@triton.jit
-def _convert_req_index_to_global_index_kernel(
-    req_id_ptr,  # int32 [num_tokens]
-    block_table_ptr,  # int32 [num_requests, max_num_blocks_per_req]
-    token_indices_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
-    out_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
-    valid_count_ptr,  # int32 [num_tokens] - output valid count per row
-    prefill_request_id_ptr,  # int32 [num_tokens], -1 for decode, >=0 for prefill
-    workspace_starts_ptr,  # int32 [num_prefill_reqs+1] or nullptr
-    # shapes (compile-time where possible)
-    max_num_blocks_per_req: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-    BLOCK_STRIDE_ROWS: tl.constexpr,
-    BLOCK_N: tl.constexpr,  # tile width along columns
-    HAS_PREFILL: tl.constexpr,
-    COUNT_VALID: tl.constexpr,  # whether to count valid indices
-    # BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count
-    # is an in-register reduction and needs no atomic.
-    SINGLE_TILE: tl.constexpr,
-    # When set, scatter valid slots to a contiguous prefix [0, valid_count) using
-    # valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior
-    # -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries).
-    # Requires COUNT_VALID and an out buffer pre-filled with -1. Order within the
-    # prefix is unspecified (only the selected set matters).
-    COMPACT_TO_FRONT: tl.constexpr,
-    # DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op
-    DCP_SIZE: tl.constexpr,
-    DCP_RANK: tl.constexpr,
-    DCP_INTERLEAVE: tl.constexpr,
-    # strides (in elements)
-    bt_stride0,
-    bt_stride1,
-    ti_stride0,
-    ti_stride1,
-    out_stride0,
-    out_stride1,
+
+
+class ConvertReqIndexToGlobalIndexKernel(
+    VllmTritonJitKernel["ConvertReqIndexToGlobalIndexKernel.CompileKey"]
 ):
-    # program_id(0) -> token_id (row)
-    # program_id(1) -> tile index along columns
-    token_id = tl.program_id(0)
-    tile_id = tl.program_id(1)
+    @dataclass(frozen=True)
+    class CompileKey:
+        block_size: int
+        block_stride_rows: int
+        block_n: int
+        has_prefill_workspace: bool
+        count_valid: bool
+        single_tile: bool
+        compact_to_front: bool
+        dcp_size: int
+        dcp_rank: int
+        dcp_interleave: int
+        num_warps: int
+        num_topk_tokens: int
+        block_table_stride: int
 
-    # Each program covers BLOCK_N consecutive columns
-    indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N)
+    @staticmethod
+    @triton.jit(do_not_specialize=["max_num_blocks_per_req"])
+    def kernel(
+        req_id_ptr,  # int32 [num_tokens]
+        block_table_ptr,  # int32 [num_requests, max_num_blocks_per_req]
+        token_indices_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
+        out_ptr,  # int32 [num_tokens, NUM_TOPK_TOKENS]
+        valid_count_ptr,  # int32 [num_tokens] - output valid count per row
+        prefill_request_id_ptr,  # int32 [num_tokens], -1 for decode, >=0 for prefill
+        workspace_starts_ptr,  # int32 [num_prefill_reqs+1] or nullptr
+        # shapes (compile-time where possible)
+        max_num_blocks_per_req,
+        BLOCK_SIZE: tl.constexpr,
+        BLOCK_STRIDE_ROWS: tl.constexpr,
+        BLOCK_N: tl.constexpr,  # tile width along columns
+        HAS_PREFILL: tl.constexpr,
+        COUNT_VALID: tl.constexpr,  # whether to count valid indices
+        # BLOCK_N == NUM_TOPK_TOKENS: one program owns the row, so the valid count
+        # is an in-register reduction and needs no atomic.
+        SINGLE_TILE: tl.constexpr,
+        # When set, scatter valid slots to a contiguous prefix [0, valid_count) using
+        # valid_count_ptr as an atomic slot allocator (DCP filtering leaves interior
+        # -1 gaps; the trtllm-gen sparse kernel reads the first valid_count entries).
+        # Requires COUNT_VALID and an out buffer pre-filled with -1. Order within the
+        # prefix is unspecified (only the selected set matters).
+        COMPACT_TO_FRONT: tl.constexpr,
+        # DCP de-interleave: with DCP_SIZE == 1 these are an exact no-op
+        DCP_SIZE: tl.constexpr,
+        DCP_RANK: tl.constexpr,
+        DCP_INTERLEAVE: tl.constexpr,
+        # strides (in elements)
+        bt_stride0,
+        bt_stride1,
+        ti_stride0,
+        ti_stride1,
+        out_stride0,
+        out_stride1,
+    ):
+        # program_id(0) -> token_id (row)
+        # program_id(1) -> tile index along columns
+        token_id = tl.program_id(0)
+        tile_id = tl.program_id(1)
 
-    # Load request id for this token (no mask: grid is exact)
-    req = tl.load(req_id_ptr + token_id)
+        # Each program covers BLOCK_N consecutive columns
+        indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N)
 
-    # Load token indices for this tile
-    ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1
-    tok = tl.load(ti_ptr)  # int32
+        # Load request id for this token (no mask: grid is exact)
+        req = tl.load(req_id_ptr + token_id)
 
-    # Only token == -1 should propagate as -1
-    is_invalid_tok = tok < 0
-    is_prefill = False
-    if HAS_PREFILL:
-        prefill_req_id = tl.load(prefill_request_id_ptr + token_id)
-        is_prefill = prefill_req_id >= 0
+        # Load token indices for this tile
+        ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1
+        tok = tl.load(ti_ptr)  # int32
 
-    # DCP de-interleave the global token id into this rank's local slot.
-    # Tokens are interleaved in groups of DCP_INTERLEAVE across ranks. With
-    # DCP_SIZE == 1 (and any interleave) owning_rank == 0 == DCP_RANK (never
-    # remote) and local_idx == tok, so this reduces to the non-DCP path; with
-    # DCP_INTERLEAVE == 1 it reduces to plain round-robin (tok % / // DCP_SIZE).
-    owning_rank = (tok // DCP_INTERLEAVE) % DCP_SIZE
-    is_remote = owning_rank != DCP_RANK
-    local_idx = (
-        tok // (DCP_SIZE * DCP_INTERLEAVE)
-    ) * DCP_INTERLEAVE + tok % DCP_INTERLEAVE
+        # Only token == -1 should propagate as -1
+        is_invalid_tok = tok < 0
+        is_prefill = False
+        if HAS_PREFILL:
+            prefill_req_id = tl.load(prefill_request_id_ptr + token_id)
+            is_prefill = prefill_req_id >= 0
 
-    # Compute block id and in-block offset
-    block_id = local_idx // BLOCK_SIZE
-    inblock_off = local_idx % BLOCK_SIZE
+        # DCP de-interleave the global token id into this rank's local slot.
+        # Tokens are interleaved in groups of DCP_INTERLEAVE across ranks. With
+        # DCP_SIZE == 1 (and any interleave) owning_rank == 0 == DCP_RANK (never
+        # remote) and local_idx == tok, so this reduces to the non-DCP path; with
+        # DCP_INTERLEAVE == 1 it reduces to plain round-robin (tok % / // DCP_SIZE).
+        owning_rank = (tok // DCP_INTERLEAVE) % DCP_SIZE
+        is_remote = owning_rank != DCP_RANK
+        local_idx = (
+            tok // (DCP_SIZE * DCP_INTERLEAVE)
+        ) * DCP_INTERLEAVE + tok % DCP_INTERLEAVE
 
-    # Guard block_table access
-    valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)
-    bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
-    is_invalid_tok |= ~valid_block | is_remote
-    base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0)
-    out_val = base * BLOCK_STRIDE_ROWS + inblock_off
+        # Compute block id and in-block offset
+        block_id = local_idx // BLOCK_SIZE
+        inblock_off = local_idx % BLOCK_SIZE
 
-    # Override with prefill output if prefill is enabled
-    if HAS_PREFILL:
-        workspace_start = tl.load(
-            workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0
-        )
-        prefill_out = workspace_start + tok
-        out_val = tl.where(is_prefill, prefill_out, out_val)
-    out_val = tl.where(is_invalid_tok, -1, out_val)
+        # Guard block_table access
+        valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)
+        bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1
+        is_invalid_tok |= ~valid_block | is_remote
+        base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0)
+        out_val = base * BLOCK_STRIDE_ROWS + inblock_off
 
-    if COMPACT_TO_FRONT:
-        # Scatter valid slots to a contiguous prefix. A per-tile exclusive prefix
-        # sum gives each valid lane a distinct local offset; one atomic add of the
-        # tile's valid count reserves a contiguous base across racing tiles. The
-        # out buffer is pre-filled with -1, so unwritten tail slots stay -1.
-        # With no racing tiles the base is 0 and the allocator becomes a store.
-        is_valid = (~is_invalid_tok).to(tl.int32)
-        local_offset = tl.cumsum(is_valid) - is_valid
-        tile_valid_count = tl.sum(is_valid)
-        if SINGLE_TILE:
-            base = 0
-            tl.store(valid_count_ptr + token_id, tile_valid_count)
-        else:
-            base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
-        dest = base + local_offset
-        out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1
-        tl.store(out_ptr_dest, out_val, mask=is_valid == 1)
-    else:
-        # Store results in place (input column == output column).
-        out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1
-        tl.store(out_ptr_ij, out_val)
+        # Override with prefill output if prefill is enabled
+        if HAS_PREFILL:
+            workspace_start = tl.load(
+                workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0
+            )
+            prefill_out = workspace_start + tok
+            out_val = tl.where(is_prefill, prefill_out, out_val)
+        out_val = tl.where(is_invalid_tok, -1, out_val)
 
-        # Accumulate the tile's valid count into the row total; a single tile's
-        # reduction *is* the total.
-        if COUNT_VALID:
-            tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32))
+        if COMPACT_TO_FRONT:
+            # Scatter valid slots to a contiguous prefix. A per-tile exclusive prefix
+            # sum gives each valid lane a distinct local offset; one atomic add of the
+            # tile's valid count reserves a contiguous base across racing tiles. The
+            # out buffer is pre-filled with -1, so unwritten tail slots stay -1.
+            # With no racing tiles the base is 0 and the allocator becomes a store.
+            is_valid = (~is_invalid_tok).to(tl.int32)
+            local_offset = tl.cumsum(is_valid) - is_valid
+            tile_valid_count = tl.sum(is_valid)
             if SINGLE_TILE:
+                base = 0
                 tl.store(valid_count_ptr + token_id, tile_valid_count)
             else:
-                tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
+                base = tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
+            dest = base + local_offset
+            out_ptr_dest = out_ptr + token_id * out_stride0 + dest * out_stride1
+            tl.store(out_ptr_dest, out_val, mask=is_valid == 1)
+        else:
+            # Store results in place (input column == output column).
+            out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1
+            tl.store(out_ptr_ij, out_val)
+
+            # Accumulate the tile's valid count into the row total; a single tile's
+            # reduction *is* the total.
+            if COUNT_VALID:
+                tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32))
+                if SINGLE_TILE:
+                    tl.store(valid_count_ptr + token_id, tile_valid_count)
+                else:
+                    tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)
+
+    def dispatch(  # type: ignore[override]
+        self,
+        *,
+        BLOCK_SIZE: int,
+        BLOCK_STRIDE_ROWS: int,
+        BLOCK_N: int,
+        NUM_TOPK_TOKENS: int,
+        HAS_PREFILL_WORKSPACE: bool,
+        COUNT_VALID: bool,
+        COMPACT_TO_FRONT: bool,
+        DCP_SIZE: int,
+        DCP_RANK: int,
+        DCP_INTERLEAVE: int,
+        max_num_blocks_per_req: int,
+    ) -> CompileKey:
+        tiling = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, COUNT_VALID)
+        return self.CompileKey(
+            block_size=BLOCK_SIZE,
+            block_stride_rows=BLOCK_STRIDE_ROWS,
+            block_n=tiling[1],
+            has_prefill_workspace=HAS_PREFILL_WORKSPACE,
+            count_valid=COUNT_VALID,
+            single_tile=tiling[0],
+            compact_to_front=COMPACT_TO_FRONT,
+            dcp_size=DCP_SIZE,
+            dcp_rank=DCP_RANK,
+            dcp_interleave=DCP_INTERLEAVE,
+            num_warps=tiling[3],
+            num_topk_tokens=NUM_TOPK_TOKENS,
+            block_table_stride=triton_scalar_specialization_rep(max_num_blocks_per_req),
+        )
+
+    def get_warmup_keys(
+        self,
+        vllm_config: VllmConfig,
+        *,
+        block_stride_rows: int,
+    ) -> list[CompileKey]:
+        block_size = vllm_config.cache_config.block_size
+        dcp_size = vllm_config.parallel_config.decode_context_parallel_size
+        dcp_interleave = vllm_config.parallel_config.cp_kv_cache_interleave_size
+        dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0
+        num_topk_tokens = vllm_config.model_config.hf_config.index_topk
+        max_num_blocks = cdiv(
+            vllm_config.model_config.max_model_len,
+            block_size * dcp_size,
+        )
+        max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size)
+        return self._trace_dispatch(self.dispatch)(
+            zip_inputs(
+                dict(
+                    HAS_PREFILL_WORKSPACE=False,
+                    COUNT_VALID=False,
+                    COMPACT_TO_FRONT=False,
+                    DCP_SIZE=1,
+                    DCP_RANK=0,
+                    DCP_INTERLEAVE=1,
+                ),
+                dict(
+                    HAS_PREFILL_WORKSPACE=False,
+                    COUNT_VALID=True,
+                    COMPACT_TO_FRONT=False,
+                    DCP_SIZE=1,
+                    DCP_RANK=0,
+                    DCP_INTERLEAVE=1,
+                ),
+                dict(
+                    HAS_PREFILL_WORKSPACE=True,
+                    COUNT_VALID=True,
+                    COMPACT_TO_FRONT=False,
+                    DCP_SIZE=1,
+                    DCP_RANK=0,
+                    DCP_INTERLEAVE=1,
+                ),
+                dict(
+                    HAS_PREFILL_WORKSPACE=False,
+                    COUNT_VALID=True,
+                    COMPACT_TO_FRONT=False,
+                    DCP_SIZE=dcp_size,
+                    DCP_RANK=dcp_rank,
+                    DCP_INTERLEAVE=dcp_interleave,
+                ),
+                dict(
+                    HAS_PREFILL_WORKSPACE=False,
+                    COUNT_VALID=True,
+                    COMPACT_TO_FRONT=dcp_size > 1,
+                    DCP_SIZE=dcp_size,
+                    DCP_RANK=dcp_rank,
+                    DCP_INTERLEAVE=dcp_interleave,
+                ),
+            ),
+            BLOCK_SIZE=block_size,
+            BLOCK_STRIDE_ROWS=block_stride_rows,
+            BLOCK_N=128,
+            NUM_TOPK_TOKENS=num_topk_tokens,
+            max_num_blocks_per_req=max_num_blocks_per_req,
+        )
+
+    def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]:
+        int32_ptr = TritonWarmupTensor(torch.int32)
+        token_shape = (1, compile_key.num_topk_tokens)
+        token_strides = (16, 1)
+        return dict(
+            req_id=int32_ptr,
+            block_table=TritonWarmupTensor(
+                torch.int32,
+                shape=(1, compile_key.block_table_stride),
+            ),
+            token_indices=TritonWarmupTensor(
+                torch.int32,
+                shape=token_shape,
+                strides=token_strides,
+            ),
+            out=TritonWarmupTensor(
+                torch.int32,
+                shape=token_shape,
+                strides=token_strides,
+            ),
+            valid_counts=int32_ptr if compile_key.count_valid else None,
+            prefill_workspace_request_ids=(
+                int32_ptr if compile_key.has_prefill_workspace else None
+            ),
+            prefill_workspace_starts=(
+                int32_ptr if compile_key.has_prefill_workspace else None
+            ),
+            max_num_blocks_per_req=1,
+            BLOCK_SIZE=compile_key.block_size,
+            BLOCK_STRIDE_ROWS=compile_key.block_stride_rows,
+            BLOCK_N=compile_key.block_n,
+            NUM_TOPK_TOKENS=compile_key.num_topk_tokens,
+            HAS_PREFILL_WORKSPACE=compile_key.has_prefill_workspace,
+            COUNT_VALID=compile_key.count_valid,
+            COMPACT_TO_FRONT=compile_key.compact_to_front,
+            DCP_SIZE=compile_key.dcp_size,
+            DCP_RANK=compile_key.dcp_rank,
+            DCP_INTERLEAVE=compile_key.dcp_interleave,
+        )
+
+    @kernel_launcher
+    def __call__(
+        self,
+        req_id: torch.Tensor,
+        block_table: torch.Tensor,
+        token_indices: torch.Tensor,
+        out: torch.Tensor,
+        valid_counts: torch.Tensor | None,
+        prefill_workspace_request_ids: torch.Tensor | None,
+        prefill_workspace_starts: torch.Tensor | None,
+        *,
+        max_num_blocks_per_req: int,
+        BLOCK_SIZE: int,
+        BLOCK_STRIDE_ROWS: int,
+        BLOCK_N: int,
+        NUM_TOPK_TOKENS: int,
+        HAS_PREFILL_WORKSPACE: bool,
+        COUNT_VALID: bool,
+        COMPACT_TO_FRONT: bool,
+        DCP_SIZE: int,
+        DCP_RANK: int,
+        DCP_INTERLEAVE: int,
+    ) -> LaunchSpec:
+        single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(
+            NUM_TOPK_TOKENS, BLOCK_N, COUNT_VALID
+        )
+        # Exact 2D grid: tokens × column tiles
+        return (req_id.shape[0], tiles_per_row), dict(
+            valid_count_ptr=valid_counts,
+            prefill_request_id_ptr=prefill_workspace_request_ids,
+            workspace_starts_ptr=prefill_workspace_starts,
+            BLOCK_N=block_n,
+            HAS_PREFILL=HAS_PREFILL_WORKSPACE,
+            SINGLE_TILE=single_tile,
+            # Strides in elements
+            bt_stride0=block_table.stride(0),
+            bt_stride1=block_table.stride(1),
+            ti_stride0=token_indices.stride(0),
+            ti_stride1=token_indices.stride(1),
+            out_stride0=out.stride(0),
+            out_stride1=out.stride(1),
+            num_warps=num_warps,
+        )
 
 
 def _remap_tiling(
@@ -231,9 +452,7 @@ def triton_convert_req_index_to_global_index(
     num_tokens = req_id.shape[0]
     max_num_blocks_per_req = block_table.shape[1]
 
-    single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(
-        NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts
-    )
+    single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, return_valid_counts)
 
     # Ensure contiguous tensors on the same device
     req_id_c = req_id.contiguous()
@@ -247,11 +466,6 @@ def triton_convert_req_index_to_global_index(
         alloc = torch.empty if single_tile else torch.zeros
         valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)
 
-    # Strides in elements
-    bt_stride0, bt_stride1 = block_table_c.stride()
-    ti_stride0, ti_stride1 = token_indices_c.stride()
-    out_stride0, out_stride1 = out.stride()
-
     # Prepare prefill pointers
     if HAS_PREFILL_WORKSPACE:
         assert prefill_workspace_request_ids is not None  # for mypy
@@ -259,10 +473,7 @@ def triton_convert_req_index_to_global_index(
         assert prefill_workspace_request_ids.is_contiguous()
         assert prefill_workspace_starts.is_contiguous()
 
-    # Exact 2D grid: tokens × column tiles
-    grid = (num_tokens, tiles_per_row)
-
-    _convert_req_index_to_global_index_kernel[grid](
+    _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL(
         req_id_c,
         block_table_c,
         token_indices_c,
@@ -271,26 +482,20 @@ def triton_convert_req_index_to_global_index(
         prefill_workspace_request_ids,
         prefill_workspace_starts,
         # shapes / constexprs
-        max_num_blocks_per_req,
-        BLOCK_SIZE,
-        BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE,
-        block_n,
-        HAS_PREFILL_WORKSPACE,
-        return_valid_counts,
-        single_tile,
-        False,  # COMPACT_TO_FRONT: keep input column == output column
+        max_num_blocks_per_req=max_num_blocks_per_req,
+        BLOCK_SIZE=BLOCK_SIZE,
+        BLOCK_STRIDE_ROWS=(
+            BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE
+        ),
+        BLOCK_N=BLOCK_N,
+        NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
+        HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE,
+        COUNT_VALID=return_valid_counts,
+        COMPACT_TO_FRONT=False,
         # DCP disabled (no-op de-interleave)
-        1,
-        0,
-        1,
-        # strides
-        bt_stride0,
-        bt_stride1,
-        ti_stride0,
-        ti_stride1,
-        out_stride0,
-        out_stride1,
-        num_warps=num_warps,
+        DCP_SIZE=1,
+        DCP_RANK=0,
+        DCP_INTERLEAVE=1,
     )
 
     if return_valid_counts:
@@ -362,9 +567,7 @@ def triton_filter_and_convert_dcp_index(
     count_valid = return_valid_counts or compact_valid_to_front
 
     # The compaction builds on the counting, so it shares the tiling.
-    single_tile, block_n, tiles_per_row, num_warps = _remap_tiling(
-        NUM_TOPK_TOKENS, BLOCK_N, count_valid
-    )
+    single_tile, _, _, _ = _remap_tiling(NUM_TOPK_TOKENS, BLOCK_N, count_valid)
 
     if compact_valid_to_front:
         out = torch.full_like(token_indices_c, -1)
@@ -377,11 +580,7 @@ def triton_filter_and_convert_dcp_index(
         alloc = torch.empty if single_tile else torch.zeros
         valid_counts = alloc(num_tokens, dtype=torch.int32, device=token_indices.device)
 
-    bt_stride0, bt_stride1 = block_table_c.stride()
-    ti_stride0, ti_stride1 = token_indices_c.stride()
-    out_stride0, out_stride1 = out.stride()
-
-    _convert_req_index_to_global_index_kernel[(num_tokens, tiles_per_row)](
+    _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL(
         req_id_c,
         block_table_c,
         token_indices_c,
@@ -390,27 +589,25 @@ def triton_filter_and_convert_dcp_index(
         # No prefill workspace on the DCP decode path.
         None,
         None,
-        max_num_blocks_per_req,
-        BLOCK_SIZE,
-        BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE,
-        block_n,
-        False,  # HAS_PREFILL
-        count_valid,
-        single_tile,
-        compact_valid_to_front,
-        dcp_size,
-        dcp_rank,
-        cp_kv_cache_interleave_size,
-        bt_stride0,
-        bt_stride1,
-        ti_stride0,
-        ti_stride1,
-        out_stride0,
-        out_stride1,
-        num_warps=num_warps,
+        max_num_blocks_per_req=max_num_blocks_per_req,
+        BLOCK_SIZE=BLOCK_SIZE,
+        BLOCK_STRIDE_ROWS=(
+            BLOCK_STRIDE_ROWS if BLOCK_STRIDE_ROWS is not None else BLOCK_SIZE
+        ),
+        BLOCK_N=BLOCK_N,
+        NUM_TOPK_TOKENS=NUM_TOPK_TOKENS,
+        HAS_PREFILL_WORKSPACE=False,
+        COUNT_VALID=count_valid,
+        COMPACT_TO_FRONT=compact_valid_to_front,
+        DCP_SIZE=dcp_size,
+        DCP_RANK=dcp_rank,
+        DCP_INTERLEAVE=cp_kv_cache_interleave_size,
     )
 
     if return_valid_counts:
         assert valid_counts is not None
         return out, valid_counts
     return out
+
+
+_CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL = ConvertReqIndexToGlobalIndexKernel()
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index a7db9a85b4..3886b28a7a 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -382,9 +382,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             model_loader = get_model_loader(self.vllm_config.load_config)
             logger.info_once("Loading model from scratch...")
 
-            self.model = model_loader.load_model(
-                vllm_config=self.vllm_config, model_config=self.vllm_config.model_config
-            )
+            # Capture warmup providers selected while constructing the model.
+            with self.jit_warmup_registry.activate():
+                self.model = model_loader.load_model(
+                    vllm_config=self.vllm_config,
+                    model_config=self.vllm_config.model_config,
+                )
             if self.lora_config:
                 self.model = self.load_lora_model(
                     self.model, self.vllm_config, self.device
@@ -668,15 +671,17 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.speculator.init_cudagraph_manager(cudagraph_mode)
 
         self.kv_caches: list[torch.Tensor] = []
-        kv_caches_dict = init_kv_cache(
-            self.kv_caches,
-            self.compilation_config.static_forward_context,
-            self.kv_cache_config,
-            self.device,
-            self.kernel_block_sizes,
-            self.vllm_config,
-            kv_cache_allocation_context=kv_cache_allocation_context,
-        )
+        # Capture warmup providers that depend on allocated KV-cache strides.
+        with self.jit_warmup_registry.activate():
+            kv_caches_dict = init_kv_cache(
+                self.kv_caches,
+                self.compilation_config.static_forward_context,
+                self.kv_cache_config,
+                self.device,
+                self.kernel_block_sizes,
+                self.vllm_config,
+                kv_cache_allocation_context=kv_cache_allocation_context,
+            )
         if is_profiling:
             self.kv_connector = NO_OP_KV_CONNECTOR
         else:
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index fdab24df0e..84c87fb79c 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -5393,9 +5393,11 @@ class GPUModelRunner(
                 if load_dummy_weights:
                     self.load_config.load_format = "dummy"
                 model_loader = get_model_loader(self.load_config)
-                self.model = model_loader.load_model(
-                    vllm_config=self.vllm_config, model_config=self.model_config
-                )
+                # Capture warmup providers selected while constructing the model.
+                with self.jit_warmup_registry.activate():
+                    self.model = model_loader.load_model(
+                        vllm_config=self.vllm_config, model_config=self.model_config
+                    )
                 if self.lora_config:
                     self.model = self.load_lora_model(
                         self.model, self.vllm_config, self.device
@@ -7500,11 +7502,13 @@ class GPUModelRunner(
 
         # Reinitialize need to after initialize_attn_backend
         self.may_reinitialize_input_batch(kv_cache_config, kernel_block_sizes)
-        kv_caches = self.initialize_kv_cache_tensors(
-            kv_cache_config,
-            kernel_block_sizes,
-            kv_cache_allocation_context=kv_cache_allocation_context,
-        )
+        # Capture warmup providers that depend on allocated KV-cache strides.
+        with self.jit_warmup_registry.activate():
+            kv_caches = self.initialize_kv_cache_tensors(
+                kv_cache_config,
+                kernel_block_sizes,
+                kv_cache_allocation_context=kv_cache_allocation_context,
+            )
 
         if (
             self.speculative_config