[Platform] Thread vendor KV pool classes through the in-tree builders; add PlatformCapabilities

docs/docs/hardware-platforms/plugin.mdx

@@ -303,15 +303,21 @@ class MyDeviceMixin(DeviceMixin):
 **5. `platform.py`** — SRT platform:
 
 ```python
-from sglang.srt.platforms.interface import SRTPlatform
+from sglang.srt.platforms.interface import PlatformCapabilities, SRTPlatform
 from my_platform_plugin.device import MyDeviceMixin
 
 class MySRTPlatform(SRTPlatform, MyDeviceMixin):
+    capabilities = PlatformCapabilities(supports_triton=False, graph_capture=True)
+
     def get_default_attention_backend(self) -> str: ...
-    def support_cuda_graph(self) -> bool: ...
-    # ... override other methods as needed
+    # ... override other methods as needed; vendor-private helpers start with `_`
 ```
 
+Every public name on the subclass (and on any mixin it composes) must exist on
+`SRTPlatform` with a compatible signature. The check runs when the class is
+defined, so a hook that core renamed or never had raises `TypeError` naming it
+instead of silently never being called.
+
 **6. Install and verify:**
 
 ```bash
@@ -535,7 +541,11 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
   </tbody>
 </table>
 
-#### Capability Flags (from SRTPlatform)
+#### Capabilities (from SRTPlatform)
+
+`capabilities` is a frozen `PlatformCapabilities` class attribute. Core reads it
+before constructing anything and picks a code path the platform can serve; every
+field defaults to the conservative answer.
 
 <table style={{width: "100%", borderCollapse: "collapse", tableLayout: "fixed"}}>
   <colgroup>
@@ -545,26 +555,31 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
   </colgroup>
   <thead>
     <tr>
-      <th>Method</th>
+      <th>Field</th>
       <th>Default</th>
-      <th>Description</th>
+      <th>What core does with it</th>
     </tr>
   </thead>
   <tbody>
     <tr>
-      <td><code>support_cuda_graph()</code></td>
+      <td><code>supports_triton</code></td>
       <td><code>False</code></td>
-      <td>Whether device graph capture is supported (plain CUDA graph)</td>
+      <td>When False, the paged allocator and the req-to-token writers use their torch-native paths instead of Triton kernels.</td>
     </tr>
     <tr>
-      <td><code>support_piecewise_cuda_graph()</code></td>
+      <td><code>graph_capture</code></td>
       <td><code>False</code></td>
-      <td>Whether piecewise CUDA graph (torch.compile backend) is supported</td>
+      <td>Runs the decode graph runner (<code>get_graph_runner_cls()</code>).</td>
     </tr>
     <tr>
-      <td><code>supports_fp8()</code></td>
+      <td><code>piecewise_graph</code></td>
       <td><code>False</code></td>
-      <td>Whether FP8 quantization is supported</td>
+      <td>Runs the prefill piecewise compilation backend.</td>
+    </tr>
+    <tr>
+      <td><code>hicache_device_kernels</code></td>
+      <td><code>False</code></td>
+      <td>The sgl_kernel HiCache transfer / write-back kernels are available; otherwise host pools skip their staging buffers.</td>
     </tr>
   </tbody>
 </table>
@@ -592,37 +607,27 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
     </tr>
     <tr>
       <td><code>get_graph_runner_cls()</code></td>
-      <td><code>raise NotImplementedError</code></td>
+      <td><code>return None</code></td>
       <td>Graph Runner class</td>
     </tr>
     <tr>
-      <td><code>get_mha_kv_pool_cls()</code></td>
-      <td><code>raise NotImplementedError</code></td>
-      <td>MHA KV cache pool class</td>
-    </tr>
-    <tr>
-      <td><code>get_mla_kv_pool_cls()</code></td>
-      <td><code>raise NotImplementedError</code></td>
-      <td>MLA KV cache pool class</td>
-    </tr>
-    <tr>
-      <td><code>get_dsa_kv_pool_cls()</code></td>
-      <td><code>raise NotImplementedError</code></td>
-      <td>DSA KV cache pool class (DeepSeek V3.2)</td>
+      <td><code>get_kv_pool_cls(kind=...)</code></td>
+      <td><code>return None</code></td>
+      <td>Return this platform&#39;s KV pool class for <code>kind</code> in <code>"mha"</code> / <code>"mla"</code> / <code>"dsa"</code>, or <code>None</code> for the in-tree class. The class must subclass <code>MHATokenToKVPool</code> / <code>MLATokenToKVPool</code> / <code>DSATokenToKVPool</code>; core constructs it with exactly the keyword arguments the in-tree class receives, at every site the in-tree class is built (standalone, inside the SWA composite, as the hybrid-linear full-attention leaf). Override <code>_create_buffers</code> to change how the backing storage is carved -- addressing, index translation and PD registration stay core&#39;s. If you override <code>_create_buffers</code> entirely you also own <code>_kv_buffer_descs</code> (PD transfer and post-capture backing read it); override <code>_create_buffers_normal</code> or <code>_kv_buffer_shapes</code> instead to keep the base descriptors, which require the leading buffer axis to be tokens or pages. The in-tree pool allocates on <code>device</code> and is correct on any torch device, so most platforms return <code>None</code>.</td>
     </tr>
     <tr>
       <td><code>get_paged_allocator_cls()</code></td>
-      <td><code>raise NotImplementedError</code></td>
-      <td>Paged allocator class</td>
+      <td><code>return None</code></td>
+      <td>Paged allocator class, honored at every paged-allocator construction site. With <code>capabilities.supports_triton=False</code> the in-tree allocator already uses torch-native kernels, so most platforms return <code>None</code>. At <code>page_size == 1</code> core uses the flat allocator on every platform.</td>
     </tr>
     <tr>
       <td><code>get_quantization_config(quantization)</code></td>
-      <td><code>raise NotImplementedError</code></td>
+      <td><code>return None</code></td>
       <td>Return hardware-specific quantization config for the specific quantization scheme, raise an error if not supported or return None to use the default config.</td>
     </tr>
     <tr>
       <td><code>get_piecewise_backend_cls()</code></td>
-      <td><code>raise NotImplementedError</code></td>
+      <td><code>return None</code></td>
       <td>Piecewise compilation backend class</td>
     </tr>
     <tr>
@@ -664,6 +669,11 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
       <td>In each worker, before model construction</td>
       <td>One-time backend initialization</td>
     </tr>
+    <tr>
+      <td><code>post_load_model(model)</code></td>
+      <td>In each worker, once the model&#39;s weights are loaded</td>
+      <td>Relocate parameters, swap quant methods or install forward hooks on an in-tree model class; mutate <code>model</code> in place</td>
+    </tr>
   </tbody>
 </table>
 

python/sglang/kernels/ops/kvcache/mla_buffer.py

@@ -117,6 +117,48 @@ def set_mla_kv_buffer_kernel_norope(
         tl.extra.cuda.gdc_launch_dependents()
 
 
+def set_mla_kv_buffer_naive(
+    kv_buffer: torch.Tensor,
+    loc: torch.Tensor,
+    cache_k_nope: torch.Tensor,
+    cache_k_rope: Optional[torch.Tensor],
+    *,
+    reserved_skip_index: int,
+    dcp_world_size: int,
+    dcp_rank: int,
+) -> None:
+    """Torch-native scatter for platforms that cannot launch Triton.
+
+    Mirrors ``set_mla_kv_buffer_kernel``: rows land at
+    ``kv_buffer[loc // dcp_world_size]`` as nope then rope, a loc equal to
+    ``reserved_skip_index`` is dropped, and under DCP a rank keeps only the
+    locs it owns. Duplicate locs race here exactly as they do in the kernel.
+    """
+    # view, not reshape: the destination must alias kv_buffer, and a layout
+    # reshape would have to copy is one this write cannot serve. The sources
+    # are read-only, so a copy there is harmless.
+    rows = kv_buffer.view(kv_buffer.shape[0], -1)
+    nope = cache_k_nope.reshape(cache_k_nope.shape[0], -1)
+    nope_dim = nope.shape[-1]
+    has_rope = cache_k_rope is not None and cache_k_rope.numel() > 0
+
+    loc = loc.to(torch.int64)
+    if has_rope:
+        is_valid = loc != reserved_skip_index
+    else:
+        # set_mla_kv_buffer_kernel_norope takes no skip index and writes every
+        # loc; keep the two paths bit-identical rather than fixing it here.
+        is_valid = torch.ones_like(loc, dtype=torch.bool)
+    if dcp_world_size > 1:
+        is_valid &= loc % dcp_world_size == dcp_rank
+    dst = loc[is_valid] // dcp_world_size
+
+    rows[dst, :nope_dim] = nope[is_valid]
+    if has_rope:
+        rope = cache_k_rope.reshape(cache_k_rope.shape[0], -1)
+        rows[dst, nope_dim : nope_dim + rope.shape[-1]] = rope[is_valid]
+
+
 def _set_mla_kv_buffer_impl(
     kv_buffer: torch.Tensor,
     loc: torch.Tensor,
@@ -139,6 +181,8 @@ def _set_mla_kv_buffer_impl(
     - Otherwise: Triton kernel with ``BLOCK = next_pow2(nope_dim + rope_dim)``,
       i.e. one CTA per loc covering the entire row in one tile. This is the
       path for SM<90 and for shapes that violate the TMA 16-byte alignment.
+    - Platforms declaring ``capabilities.supports_triton=False``: the
+      torch-native scatter, which no device-specific kernel backs.
 
     Speedup vs the legacy BLOCK=128 Triton kernel on GB300 (BF16, nope=512,
     rope=64): ~1.05x at bs=8, ~1.5x at bs=128, 3.5x at bs=512, **11.7x at
@@ -153,6 +197,23 @@ def _set_mla_kv_buffer_impl(
     Shared body of the two entry points below; the owner rule reaches it as
     ``1, 0`` (nothing to select) or as the live topology.
     """
+    # Imported here, not at module scope: `current_platform` is a lazy module
+    # attribute, so a top-level `from ... import` resolves the platform during
+    # this module's import.
+    from sglang.srt.platforms import current_platform
+
+    if not current_platform.capabilities.supports_triton:
+        set_mla_kv_buffer_naive(
+            kv_buffer,
+            loc,
+            cache_k_nope,
+            cache_k_rope,
+            reserved_skip_index=reserved_skip_index,
+            dcp_world_size=dcp_world_size,
+            dcp_rank=dcp_rank,
+        )
+        return
+
     has_rope = cache_k_rope is not None and cache_k_rope.numel() > 0
     n_loc = loc.numel()
     nope_dim = cache_k_nope.shape[-1]

python/sglang/srt/arg_groups/cuda_graph_hook.py

@@ -203,7 +203,7 @@ def disable_tc_piecewise_cudagraph_if_incompatible(server_args: Any):
             "OOT platform without piecewise support",
             lambda: (
                 current_platform.is_out_of_tree()
-                and not current_platform.support_piecewise_cuda_graph()
+                and not current_platform.capabilities.piecewise_graph
             ),
         ),
         (

python/sglang/srt/compilation/backend.py

@@ -26,6 +26,7 @@
 from sglang.srt.compilation.xpu_piecewise_backend import XPUPiecewiseBackend
 from sglang.srt.environ import envs
 from sglang.srt.platforms import current_platform
+from sglang.srt.platforms.interface import require_out_of_tree_impl
 from sglang.srt.utils.common import is_npu, is_xpu
 
 logger = logging.getLogger(__name__)
@@ -52,14 +53,19 @@ def make_backend(
     sglang_backend,
 ):
 
-    if current_platform.is_out_of_tree():
-        backend_cls = current_platform.get_piecewise_backend_cls()
-    elif is_xpu():
-        backend_cls = XPUPiecewiseBackend
-    elif is_npu():
-        backend_cls = NPUPiecewiseBackend
-    else:
-        backend_cls = CUDAPiecewiseBackend
+    backend_cls = current_platform.get_piecewise_backend_cls()
+    if backend_cls is None:
+        require_out_of_tree_impl(
+            current_platform,
+            hook="get_piecewise_backend_cls()",
+            subsystem="piecewise compilation backend",
+        )
+        if is_xpu():
+            backend_cls = XPUPiecewiseBackend
+        elif is_npu():
+            backend_cls = NPUPiecewiseBackend
+        else:
+            backend_cls = CUDAPiecewiseBackend
     return backend_cls(
         graph,
         compile_config,

python/sglang/srt/distributed/parallel_state.py

@@ -2377,8 +2377,9 @@ def get_default_distributed_backend(device: str) -> str:
     # ``from ... import current_platform``) so each call resolves through the
     # platforms package's lazy ``__getattr__`` and picks up runtime overrides
     # of ``_current_platform`` (e.g. in tests).
-    if device == platforms.current_platform.device_type:
-        return platforms.current_platform.get_torch_distributed_backend_str()
+    platform = platforms.current_platform
+    if device in (platform.device_type, platform.device_name):
+        return platform.get_torch_distributed_backend_str()
     return _DEVICE_TO_DISTRIBUTED_BACKEND.get(device, "gloo")
 
 

python/sglang/srt/kv_canary/api.py

@@ -44,7 +44,7 @@ def torch_reference_conflicts_with_decode_graph(device: torch.device) -> bool:
     # than waves through: a startup error beats a canary that reports clean forever.
     return (
         use_torch_reference(device)
-        and current_platform.support_cuda_graph()
+        and current_platform.capabilities.graph_capture
         and not check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED)
     )
 

python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py

@@ -30,7 +30,7 @@
     mqa_logits_should_chunk,
 )
 from sglang.srt.layers.layernorm import LayerNorm
-from sglang.srt.layers.utils import MultiPlatformOp
+from sglang.srt.layers.utils.multi_platform import MultiPlatformOp
 from sglang.srt.utils import add_prefix, ceil_align, is_cuda, is_hip, is_npu
 
 if is_cuda():

python/sglang/srt/layers/attention/qsa/qsa_indexer.py

@@ -27,7 +27,7 @@
 from sglang.srt.layers.layernorm import GemmaRMSNorm
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb
-from sglang.srt.layers.utils import MultiPlatformOp
+from sglang.srt.layers.utils.multi_platform import MultiPlatformOp
 from sglang.srt.model_executor.runner import get_is_capture_mode
 
 # Cap on the fp32 [query_rows, compressed_keys] prefill logits workspace;

python/sglang/srt/layers/utils/__init__.py

@@ -1,6 +1,2 @@
 # Temp workaround, make layer utils more fine-grained later
 from sglang.srt.layers.utils.common import *
-
-# Deprecated re-export kept for external users; in-repo code subclasses
-# sglang.kernels.fused_op.BaseFusedOp directly (RFC #29630).
-from sglang.srt.layers.utils.multi_platform import MultiPlatformOp

python/sglang/srt/layers/utils/multi_platform.py

@@ -8,8 +8,9 @@
 and the torch.compile enter/leave protocol.
 
 In-repo code must subclass ``BaseFusedOp`` directly. This alias exists only so
-out-of-tree platform plugins and external users keep importing from the old
-path while they migrate; it will be removed in a future release.
+out-of-tree platform plugins and external users keep importing from
+``sglang.srt.layers.utils.multi_platform`` while they migrate; it will be
+removed in a future release.
 """
 
 import warnings

python/sglang/srt/managers/data_parallel_controller.py

@@ -49,6 +49,7 @@
 from sglang.srt.observability.req_time_stats import DPControllerReqTimeStats
 from sglang.srt.observability.startup_time import aggregate_scheduler_startup_times
 from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
+from sglang.srt.plugins import load_plugins
 from sglang.srt.runtime_context import (
     get_device,
     get_disagg,
@@ -813,6 +814,7 @@ def run_data_parallel_controller_process(
     kill_itself_when_parent_died()
     parent_process = psutil.Process().parent()
 
+    load_plugins()
     # This process reads the config namespaces before spawning schedulers.
     publish(server_args, role="dp_controller")
     configure_logger(server_args)

python/sglang/srt/managers/detokenizer_manager.py

@@ -43,6 +43,7 @@
 )
 from sglang.srt.managers.multi_tokenizer_mixin import MultiHttpWorkerDetokenizerMixin
 from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread
+from sglang.srt.plugins import load_plugins
 from sglang.srt.runtime_context import (
     get_device,
     get_model,
@@ -542,6 +543,7 @@ def run_detokenizer_process(
 ):
     kill_itself_when_parent_died()
     setproctitle.setproctitle("sglang::detokenizer")
+    load_plugins()
     configure_logger(server_args)
     publish(server_args, role="detokenizer")
     parent_process = psutil.Process().parent()

python/sglang/srt/mem_cache/allocation.py

@@ -592,7 +592,7 @@ def alloc_paged_token_slots_decode(
     return out_cache_loc
 
 
-def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
+def alloc_for_decode_default(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
     """
     Allocate KV cache for decode batch and write to req_to_token_pool.
 
@@ -693,6 +693,17 @@ def assign_req_to_token_pool(
         load_offset += BLOCK_SIZE
 
 
+ALLOC_FOR_DECODE_FUNCS = defaultdict(lambda: alloc_for_decode_default)
+
+
+def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
+    return ALLOC_FOR_DECODE_FUNCS[_device_type(batch)](batch, token_per_req)
+
+
+def _device_type(batch: ScheduleBatch) -> str:
+    return torch.device(batch.device).type
+
+
 def assign_req_to_token_pool_func(
     req_pool_indices: torch.Tensor,
     req_to_token: torch.Tensor,
@@ -756,10 +767,7 @@ def alloc_for_spec_decode(
             last_loc = get_last_loc(
                 req_to_token_pool.req_to_token, req_pool_indices, cur_kv_lens
             )
-            device_type = getattr(
-                batch.device, "type", str(batch.device).split(":", 1)[0]
-            )
-            out_cache_loc = ALLOC_EXTEND_FUNCS[device_type](
+            out_cache_loc = ALLOC_EXTEND_FUNCS[_device_type(batch)](
                 tree_cache,
                 cur_kv_lens,
                 cur_kv_lens_cpu,