[DSv4] Integrate TRT-LLM DSv4 Attention for SM100/103 (#30805)

Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>

python/sglang/srt/arg_groups/deepseek_v4_hook.py

@@ -134,6 +134,47 @@ def apply_deepseek_v4_defaults(server_args: ServerArgs, model_arch: str) -> None
 
     run_post_process_pass(server_args, _deepseek_v4_kv_cache_dtype)
 
+    if cfg.dsv4_attn_backend == "trtllm":
+        from sglang.srt.utils.common import is_sm100_supported
+
+        assert cfg.device == "cuda" and is_sm100_supported(), (
+            "--dsv4-attn-backend trtllm requires an SM100/SM103 (Blackwell) GPU."
+        )
+        # The resolution pipeline materializes "auto" as fp8_e4m3 on CUDA.
+        assert cfg.kv_cache_dtype in ("auto", "fp8_e4m3"), (
+            "--dsv4-attn-backend trtllm requires kv_cache_dtype=fp8_e4m3, "
+            f"got {cfg.kv_cache_dtype}."
+        )
+        assert not cfg.enable_hisparse, (
+            "--dsv4-attn-backend trtllm does not support enable_hisparse."
+        )
+        assert not (
+            cfg.attn_cp_size > 1 or cfg.dcp_size > 1 or cfg.enable_prefill_cp
+        ), (
+            "--dsv4-attn-backend trtllm does not support context parallelism "
+            "(prefill CP, attention CP, or decode CP)."
+        )
+        # The trtllm backend stores KV in a 512-byte uniform-FP8 layout while
+        # FlashMLA uses the 584-byte packed layout; the PD handshake only
+        # compares kv_cache_dtype, so mismatched prefill/decode backends would
+        # pass the check and transfer garbage. Reject until the handshake
+        # carries a layout identifier and the path is tested (#37838).
+        assert cfg.disaggregation_mode == "null", (
+            "--dsv4-attn-backend trtllm does not support PD disaggregation yet "
+            "(uniform-FP8 KV layout is not part of the PD handshake; see "
+            "https://github.com/sgl-project/sglang/issues/37838)."
+        )
+        # The trtllm-gen semaphore buffer is sized from the prefill chunk
+        # bound; with chunking disabled a single long request has no bound.
+        assert cfg.chunked_prefill_size is not None and cfg.chunked_prefill_size > 0, (
+            "--dsv4-attn-backend trtllm requires chunked prefill "
+            "(--chunked-prefill-size > 0)."
+        )
+        logger.info(
+            "DeepSeek V4 attention: trtllm backend enabled "
+            "(uniform-FP8 KV pool, decode + sparse prefill)."
+        )
+
     if cfg.max_running_requests is None:
         declare_resolution(
             server_args,

python/sglang/srt/arg_groups/fields/exec_.py

@@ -230,6 +230,19 @@ class ExecKernel:
             resolvable=True,
         ),
     ] = None
+    dsv4_attn_backend: A[
+        str,
+        Arg(
+            help="DeepSeek V4 attention backend. 'auto' (default) resolves to "
+            "'flashmla'. 'trtllm' (opt-in, SM100/SM103 with FP8 KV cache) "
+            "switches the SWA/compressed KV pools to a "
+            "uniform 512-dim FP8 layout and runs decode and sparse prefill "
+            "through the flashinfer trtllm-gen sparse MLA kernel. The backend "
+            "choice is shared by prefill and decode.",
+            choices=["auto", "flashmla", "trtllm"],
+            resolvable=True,
+        ),
+    ] = "auto"
     dsa_paged_mqa_logits_backend: A[
         str,
         Arg(

python/sglang/srt/layers/attention/attention_registry.py

@@ -181,12 +181,15 @@ def create_dsv4_backend(runner):
         )
         return DeepseekV4HipRadixBackend(runner)
     else:
-        from sglang.srt.layers.attention.deepseek_v4_backend import (
+        from sglang.srt.layers.attention.deepseek_v4_trtllm_backend import (
-            DeepseekV4AttnBackend,
+            create_deepseek_v4_attn_backend,
         )
 
-        logger.info("Using DeepseekV4AttnBackend for dsv4 attention backend (CUDA).")
+        backend = create_deepseek_v4_attn_backend(runner)
-        return DeepseekV4AttnBackend(runner)
+        logger.info(
+            f"Using {type(backend).__name__} for dsv4 attention backend (CUDA)."
+        )
+        return backend
 
 
 @register_attention_backend("triton")

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -193,6 +193,18 @@ class DSV4AttnMetadata:
     c128_page_indices: Optional[torch.Tensor] = None
     c128_topk_lengths_clamp1: Optional[torch.Tensor] = None
 
+    # Combined decode tables. Only the c4 tail and lens vary by layer.
+    trtllm_swa_lens: Optional[torch.Tensor] = None
+    trtllm_c4_indices: Optional[torch.Tensor] = None
+    trtllm_c4_lens: Optional[torch.Tensor] = None
+    trtllm_c128_indices: Optional[torch.Tensor] = None
+    trtllm_c128_lens: Optional[torch.Tensor] = None
+    # Lazy eager-prefill caches: qmeta and per-ratio combined tables.
+    trtllm_prefill_qmeta: Optional[tuple] = None
+    trtllm_prefill_swa_lens: Optional[torch.Tensor] = None
+    trtllm_prefill_c4_indices: Optional[torch.Tensor] = None
+    trtllm_prefill_c128: Optional[tuple] = None
+
     c1_flashmla_metadata: FlashMLASchedMeta = field(init=False, repr=False)
     c4_flashmla_metadata: FlashMLASchedMeta = field(init=False, repr=False)
     c128_flashmla_metadata: FlashMLASchedMeta = field(init=False, repr=False)
@@ -236,6 +248,11 @@ class DSV4AttnMetadata:
                 "c4_sparse_topk_lengths",
                 "c4_sparse_page_indices",
                 "c4_sparse_raw_indices",
+                "trtllm_swa_lens",
+                "trtllm_c4_indices",
+                "trtllm_c4_lens",
+                "trtllm_c128_indices",
+                "trtllm_c128_lens",
             ],
             assign_fields=[
                 # Recomputed by the recorded init_forward_metadata_in_graph op
@@ -244,6 +261,11 @@ class DSV4AttnMetadata:
                 "c1_flashmla_metadata",
                 "c4_flashmla_metadata",
                 "c128_flashmla_metadata",
+                # Eager-only lazy caches are assigned, not content-copied.
+                "trtllm_prefill_qmeta",
+                "trtllm_prefill_swa_lens",
+                "trtllm_prefill_c4_indices",
+                "trtllm_prefill_c128",
             ],
         )
 
@@ -261,6 +283,12 @@ class DSV4AttnMetadata:
             "c4_topk_lengths_raw",
             "c4_topk_lengths_clamp1",
             "c4_sparse_topk_lengths",
+            # Preserve graph-captured table addresses; refill c4 per layer.
+            "trtllm_swa_lens",
+            "trtllm_c4_indices",
+            "trtllm_c4_lens",
+            "trtllm_c128_indices",
+            "trtllm_c128_lens",
         ]
         reference_assign_fields = [
             "page_table",
@@ -271,6 +299,11 @@ class DSV4AttnMetadata:
             "c1_flashmla_metadata",
             "c4_flashmla_metadata",
             "c128_flashmla_metadata",
+            # Reset eager-only caches so a replay cannot reuse another shape.
+            "trtllm_prefill_qmeta",
+            "trtllm_prefill_swa_lens",
+            "trtllm_prefill_c4_indices",
+            "trtllm_prefill_c128",
         ]
         # Keep graph-captured tensor objects alive for fields that captured
         # kernels read by address; overwrite only their contents.
@@ -399,6 +432,51 @@ class DSV4AttnMetadata:
         self.c4_flashmla_metadata = _create_flashmla_metadata()
         self.c128_flashmla_metadata = _create_flashmla_metadata()
 
+    def init_trtllm_sparse_buffers(self) -> None:
+        """Build decode tables with 128 SWA columns followed by compressed KV.
+
+        Indices use -1 for invalid entries; lens include all 128 SWA slots.
+        Only the c4 tail and lens are filled per layer.
+        """
+
+        num_tokens = self.seq_lens_casual.shape[0]
+        assert self.swa_page_indices.shape == (num_tokens, SWA_WINDOW)
+
+        # VarSeq reads rows to the 64-token tile boundary. Back every live view
+        # with an aligned parent whose extra rows contain inert values.
+        n_pad = (num_tokens + 63) // 64 * 64
+
+        def _tile_padded(fill, src=None, width=None):
+            shape = (n_pad,) if width is None else (n_pad, width)
+            buf = torch.full(shape, fill, **self.cuda_int32_kwargs)
+            if src is not None:
+                buf[:num_tokens].copy_(src)
+            return buf[:num_tokens]
+
+        if n_pad != num_tokens:
+            self.seq_lens_casual = _tile_padded(1, self.seq_lens_casual)
+            self.swa_page_indices = _tile_padded(
+                -1, self.swa_page_indices, width=SWA_WINDOW
+            )
+        self.trtllm_swa_lens = _tile_padded(SWA_WINDOW)
+        if self.c4_sparse_page_indices is not None:
+            w4 = self.c4_sparse_page_indices.shape[-1]
+            assert w4 % 4 == 0, f"{w4=}"
+            # Unwritten c4 rows must remain inert until the per-layer fill.
+            self.trtllm_c4_indices = _tile_padded(-1, width=SWA_WINDOW + w4)
+            self.trtllm_c4_indices[:, :SWA_WINDOW].copy_(self.swa_page_indices)
+            self.trtllm_c4_lens = _tile_padded(SWA_WINDOW)
+        if self.c128_page_indices is not None:
+            w128 = self.c128_page_indices.shape[-1]
+            assert w128 % 4 == 0, f"{w128=}"
+            self.trtllm_c128_indices = _tile_padded(-1, width=SWA_WINDOW + w128)
+            self.trtllm_c128_indices[:, :SWA_WINDOW].copy_(self.swa_page_indices)
+            self.trtllm_c128_indices[:, SWA_WINDOW:].copy_(self.c128_page_indices)
+            self.trtllm_c128_lens = _tile_padded(
+                SWA_WINDOW,
+                (self.c128_topk_lengths_clamp1 + SWA_WINDOW).to(torch.int32),
+            )
+
 
 @dataclass
 class DSV4Metadata:
@@ -512,6 +590,7 @@ class DeepseekV4AttnBackend(
     use_captured_forward_metadata_for_breakable_cuda_graph: bool = True
     supports_ragged_verify_graph: bool = True
     needs_cpu_seq_lens: bool = False
+    trtllm_attn: bool = False
 
     def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds:
         # Breakable-graph verify rereads shared state across segments.
@@ -1051,6 +1130,10 @@ class DeepseekV4AttnBackend(
                 req_pool_indices=req_pool_indices,
             )
         )
+        if self.trtllm_attn:
+            # DP padding can expand a length-one request into nonpositive
+            # per-token lens; trtllm-gen requires them to remain at least one.
+            seq_lens_casual = seq_lens_casual.clamp(min=1)
         core_attn_metadata = self.make_core_attn_metadata(
             req_to_token=self.req_to_token,
             req_pool_indices_repeated=req_pool_indices_repeated,
@@ -1737,6 +1820,20 @@ class DeepseekV4AttnBackend(
             extra_indices = match_num_queries(extra_indices, value=-1)
             extra_topk_lengths = match_num_queries(extra_topk_lengths, value=1)
 
+            if self.trtllm_attn:
+                # The uniform-FP8 pool is readable only by trtllm-gen.
+                return self._forward_trtllm(
+                    q=q,
+                    layer=layer,
+                    compress_ratio=compress_ratio,
+                    core_attn_metadata=core_attn_metadata,
+                    forward_batch=forward_batch,
+                    attn_sink=attn_sink,
+                    swa_page_indices=swa_page_indices,
+                    extra_indices=extra_indices,
+                    extra_topk_lengths=extra_topk_lengths,
+                )
+
             if q.ndim == 3:
                 q = q.unsqueeze(1)
             if swa_page_indices.ndim == 2:
@@ -2263,6 +2360,8 @@ class DeepseekV4AttnBackend(
         if need_compress:
             core_attn_metadata.init_compression_metadata(num_tokens)
             core_attn_metadata.init_flashmla_related(is_prefill=is_prefill)
+            if self.trtllm_attn:
+                core_attn_metadata.init_trtllm_sparse_buffers()
         else:
             core_attn_metadata.c4_sparse_topk_lengths = None
             core_attn_metadata.c4_sparse_page_indices = None
@@ -2270,6 +2369,8 @@ class DeepseekV4AttnBackend(
             core_attn_metadata.c1_flashmla_metadata = _create_flashmla_metadata()
             core_attn_metadata.c4_flashmla_metadata = None
             core_attn_metadata.c128_flashmla_metadata = None
+            if self.trtllm_attn:
+                core_attn_metadata.init_trtllm_sparse_buffers()
         return core_attn_metadata
 
     def get_dspark_swa_page_indices(
@@ -2312,14 +2413,17 @@ class DeepseekV4MultiStepBackend(DeepseekV4AttnBackend):
         self.speculative_num_steps = speculative_num_steps
         self.attn_backends: List[DeepseekV4AttnBackend] = []
         for i in range(self.speculative_num_steps):
-            self.attn_backends.append(
+            self.attn_backends.append(self._make_step_backend(model_runner, i))
-                DeepseekV4AttnBackend(
+
+    def _make_step_backend(
+        self, model_runner: ModelRunner, step_id: int
+    ) -> DeepseekV4AttnBackend:
+        return DeepseekV4AttnBackend(
             model_runner,
-                    speculative_step_id=i,
+            speculative_step_id=step_id,
             topk=self.topk,
             speculative_num_steps=self.speculative_num_steps,
         )
-            )
 
     def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None:
         for attn_backend in self.attn_backends:

python/sglang/srt/layers/attention/deepseek_v4_trtllm_backend.py

@@ -0,0 +1,532 @@
+"""DeepSeek V4 trtllm-gen sparse MLA backend for SM100/SM103.
+
+Decode and varlen prefill use a uniform 512-dim FP8 KV cache. Shared metadata
+construction preserves the base backend's CUDA-graph replay semantics.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Literal, Optional, Tuple
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.attention.deepseek_v4_backend import (
+    SWA_WINDOW,
+    DeepseekV4AttnBackend,
+    DeepseekV4MultiStepBackend,
+)
+from sglang.srt.runtime_context import (
+    get_exec,
+    get_parallel,
+    get_schedule,
+    get_spec,
+    max_prefill_buffer_tokens,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.deepseek_v4_backend import DSV4AttnMetadata
+    from sglang.srt.layers.radix_attention import RadixAttention
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+logger = logging.getLogger(__name__)
+
+# Shared zero-initialized workspace managed by the persistent-buffer lifecycle.
+_TRTLLM_GEN_WORKSPACE_SIZE_MB = 128
+
+
+def _get_trtllm_workspace_buffer(device: torch.device) -> torch.Tensor:
+    from sglang.srt.runtime_context import get_buffer
+
+    return get_buffer(
+        "trtllm_dsv4_zero_workspace",
+        lambda: torch.zeros(
+            _TRTLLM_GEN_WORKSPACE_SIZE_MB * 1024 * 1024,
+            dtype=torch.int8,
+            device=device,
+        ),
+    )
+
+
+_trtllm_semaphore_installed = False
+# Capacity is in query rows: requests x draft tokens for decode, sum_q for prefill.
+_trtllm_semaphore_rows: int = 0
+
+
+def _trtllm_query_row_capacity(model_runner: ModelRunner) -> int:
+    """Bound query rows across prefill chunks and speculative decode batches.
+
+    The DSv4 hook rejects the backend when chunked prefill is disabled, so the
+    prefill chunk bound is always finite here.
+    """
+    schedule = get_schedule()
+    rows = max(
+        schedule.max_prefill_tokens or 0,
+        max_prefill_buffer_tokens(),
+    )
+    spec = get_spec()
+    rows_per_req = (
+        (spec.speculative_num_draft_tokens or 1)
+        if spec.speculative_algorithm is not None
+        else 1
+    )
+    rows = max(rows, (schedule.max_running_requests or 0) * rows_per_req)
+    return max(rows, 1)
+
+
+def _install_persistent_trtllm_semaphores(capacity_rows: int) -> None:
+    """Install a persistent counter buffer sized by query rows.
+
+    FlashInfer sizes this private buffer by request count, while the DSV4
+    VarSeq kernel indexes it by query row. Remove this workaround once
+    FlashInfer accepts a caller-owned buffer.
+    """
+    global _trtllm_semaphore_installed, _trtllm_semaphore_rows
+    _trtllm_semaphore_rows = max(_trtllm_semaphore_rows, capacity_rows)
+    if _trtllm_semaphore_installed:
+        return
+    import flashinfer.mla._core as _fi_core
+
+    _orig = _fi_core._get_trtllm_gen_multi_ctas_kv_counter_buffer
+    # Allocate once outside graph capture. Stream ordering and the kernel's
+    # counter reset make one shared buffer safe across launches.
+    state: dict = {}
+
+    def _patched(batch_size, num_qo_heads, sm_count, device):
+        buf = state.get("buf")
+        if buf is None or buf.device != device:
+            assert not torch.cuda.is_current_stream_capturing(), (
+                "persistent trtllm semaphore buffer must be created outside "
+                "graph capture (first call is expected during eager warmup)"
+            )
+            buf = _orig(_trtllm_semaphore_rows, num_qo_heads, sm_count, device)
+            state["buf"] = buf
+        return buf
+
+    _fi_core._get_trtllm_gen_multi_ctas_kv_counter_buffer = _patched
+    _trtllm_semaphore_installed = True
+    logger.info(
+        "trtllm-gen multi-CTA semaphores: single persistent buffer sized for "
+        "%d query rows, shared across launches (flashinfer sizing WAR).",
+        _trtllm_semaphore_rows,
+    )
+
+
+def _check_trtllm_query_rows(num_rows: int) -> None:
+    # A plain exception, not assert: an over-capacity launch scribbles past
+    # the semaphore buffer, so this must fire even under python -O.
+    if num_rows > _trtllm_semaphore_rows:
+        raise RuntimeError(
+            f"trtllm-gen launch with {num_rows} query rows exceeds the persistent "
+            f"semaphore capacity of {_trtllm_semaphore_rows} rows derived from "
+            "--chunked-prefill-size / --max-prefill-tokens / "
+            "--max-running-requests; lower --chunked-prefill-size."
+        )
+
+
+class DeepseekV4TrtllmAttnBackend(DeepseekV4AttnBackend):
+    """DSV4 attention through the trtllm-gen sparse MLA kernel."""
+
+    trtllm_attn: bool = True
+
+    def __init__(
+        self,
+        model_runner: ModelRunner,
+        skip_prefill: bool = False,
+        speculative_step_id=0,
+        topk=0,
+        speculative_num_steps=0,
+    ):
+        _install_persistent_trtllm_semaphores(_trtllm_query_row_capacity(model_runner))
+        super().__init__(
+            model_runner,
+            skip_prefill=skip_prefill,
+            speculative_step_id=speculative_step_id,
+            topk=topk,
+            speculative_num_steps=speculative_num_steps,
+        )
+        assert self.token_to_kv_pool.uniform_fp8, (
+            "the trtllm backend requires the uniform-FP8 DSv4 KV pool."
+        )
+        assert not envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get(), (
+            "--dsv4-attn-backend trtllm does not support "
+            "SGLANG_OPT_USE_ONLINE_COMPRESS yet."
+        )
+        # CP round-robin reindexing breaks VarSeq's per-request query packing.
+        assert get_parallel().attn_cp_size == 1, (
+            "--dsv4-attn-backend trtllm does not support "
+            "context parallelism (attn_cp_size > 1) yet."
+        )
+        self.trtllm_workspace_buffer = _get_trtllm_workspace_buffer(self.device)
+
+    def _forward_trtllm(
+        self,
+        *,
+        q: torch.Tensor,
+        layer: RadixAttention,
+        compress_ratio: Literal[0, 4, 128],
+        core_attn_metadata: DSV4AttnMetadata,
+        forward_batch: ForwardBatch,
+        attn_sink: torch.Tensor,
+        swa_page_indices: torch.Tensor,
+        extra_indices: Optional[torch.Tensor],
+        extra_topk_lengths: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        assert attn_sink is not None
+        if (
+            forward_batch.forward_mode.is_decode_or_idle()
+            or forward_batch.forward_mode.is_target_verify()
+            or forward_batch.forward_mode.is_draft_extend_v2()
+        ):
+            return self._forward_trtllm_decode(
+                q=q,
+                layer=layer,
+                compress_ratio=compress_ratio,
+                core_attn_metadata=core_attn_metadata,
+                attn_sink=attn_sink,
+                swa_page_indices=swa_page_indices,
+                extra_indices=extra_indices,
+                extra_topk_lengths=extra_topk_lengths,
+            )
+        assert forward_batch.forward_mode.is_extend_without_speculative(), (
+            "uniform-FP8 pool cannot be read by the packed FlashMLA "
+            f"kernels; unsupported forward mode "
+            f"{forward_batch.forward_mode} under "
+            "--dsv4-attn-backend trtllm"
+        )
+        return self._forward_trtllm_prefill(
+            q=q,
+            layer=layer,
+            compress_ratio=compress_ratio,
+            forward_batch=forward_batch,
+            attn_sink=attn_sink,
+            swa_page_indices=swa_page_indices,
+            extra_indices=extra_indices,
+            extra_topk_lengths=extra_topk_lengths,
+        )
+
+    def _get_trtllm_bmm_scales(self, layer: RadixAttention) -> Tuple[float, float]:
+        """Return host scales; KV uses the store path's fixed unit scale.
+
+        Tensor scales corrupt split-KV reduction on FlashInfer < 0.6.13.
+        """
+
+        assert layer.k_scale_float is None or layer.k_scale_float == 1.0, (
+            "--dsv4-attn-backend trtllm stores KV with a "
+            "fixed per-tensor scale of 1.0; a non-unit checkpoint kv-cache "
+            f"scale (k_scale_float={layer.k_scale_float}) is not supported yet."
+        )
+        return (self.softmax_scale, 1.0)
+
+    def _trtllm_kv_cache_views(
+        self, layer_id: int, compress_ratio: Literal[0, 4, 128]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Return HND views of the uniform-FP8 SWA and compressed pools.
+
+        SWA-only layers pass the SWA pool as the required compressed tensor;
+        ``sparse_topk_lens`` masks that region.
+        """
+
+        token_to_kv_pool = self.token_to_kv_pool
+        swa_buf = token_to_kv_pool.get_swa_key_buffer_radix(layer_id)
+        swa_page_size = token_to_kv_pool.swa_kv_pool.page_size
+        swa_kv_cache = swa_buf.view(swa_buf.shape[0], 1, swa_page_size, 512)
+        if compress_ratio == 0:
+            compressed_kv_cache = swa_kv_cache
+        else:
+            extra_buf = token_to_kv_pool.get_extra_key_buffer(layer_id)
+            extra_page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
+            compressed_kv_cache = extra_buf.view(
+                extra_buf.shape[0], 1, extra_page_size, 512
+            )
+        return swa_kv_cache, compressed_kv_cache
+
+    def _forward_trtllm_decode(
+        self,
+        *,
+        q: torch.Tensor,
+        layer: RadixAttention,
+        compress_ratio: Literal[0, 4, 128],
+        core_attn_metadata: DSV4AttnMetadata,
+        attn_sink: torch.Tensor,
+        swa_page_indices: torch.Tensor,
+        extra_indices: Optional[torch.Tensor],
+        extra_topk_lengths: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        """Run sparse MLA decode with preallocated metadata tables."""
+
+        from flashinfer.mla import trtllm_batch_decode_sparse_mla_dsv4
+
+        bs, num_heads, head_dim = q.shape
+        assert head_dim == 512
+
+        # Draft-extend metadata predates DP MAX_LEN padding (#27091). Run only
+        # its covered rows and leave the discarded padding output finite.
+        n_meta_rows = core_attn_metadata.seq_lens_casual.shape[0]
+        out_pad_tail = None
+        if n_meta_rows < bs:
+            out_pad_tail = torch.zeros(
+                (bs, num_heads, 512), dtype=torch.bfloat16, device=q.device
+            )
+            q = q[:n_meta_rows]
+            swa_page_indices = swa_page_indices[:n_meta_rows]
+            if extra_indices is not None:
+                extra_indices = extra_indices[:n_meta_rows]
+            if extra_topk_lengths is not None:
+                extra_topk_lengths = extra_topk_lengths[:n_meta_rows]
+            bs = n_meta_rows
+
+        # Only the c4 tail and lens vary by layer; other table data is prebuilt.
+        assert swa_page_indices.shape == (bs, SWA_WINDOW)
+        if compress_ratio == 0:
+            # Use the metadata view backed by 64-row-aligned storage because
+            # the VarSeq kernel reads table rows to the tile boundary.
+            sparse_indices = core_attn_metadata.swa_page_indices
+            sparse_topk_lens = core_attn_metadata.trtllm_swa_lens
+        elif compress_ratio == 128:
+            sparse_indices = core_attn_metadata.trtllm_c128_indices
+            sparse_topk_lens = core_attn_metadata.trtllm_c128_lens
+        else:
+            sparse_indices = core_attn_metadata.trtllm_c4_indices
+            sparse_topk_lens = core_attn_metadata.trtllm_c4_lens
+        assert sparse_indices is not None and sparse_topk_lens is not None, (
+            "trtllm decode requires metadata built with "
+            "init_trtllm_sparse_buffers (decode-mode DSV4AttnMetadata)"
+        )
+        if sparse_indices.shape[0] != bs:
+            assert sparse_indices.shape[0] > bs, f"{sparse_indices.shape=} {bs=}"
+            sparse_indices = sparse_indices[:bs]
+        if sparse_topk_lens.shape[0] != bs:
+            assert sparse_topk_lens.shape[0] > bs, f"{sparse_topk_lens.shape=}"
+            sparse_topk_lens = sparse_topk_lens[:bs]
+
+        if compress_ratio == 4:
+            assert extra_indices is not None and extra_topk_lengths is not None
+            width = extra_indices.shape[-1]
+            assert SWA_WINDOW + width == sparse_indices.shape[1], (
+                f"{width=} {sparse_indices.shape=}"
+            )
+            sparse_indices[:, SWA_WINDOW:].copy_(extra_indices)
+            # Lens include all 128 SWA slots; seq_lens controls their validity.
+            sparse_topk_lens.copy_(extra_topk_lengths)
+            sparse_topk_lens.add_(SWA_WINDOW)
+
+        swa_kv_cache, compressed_kv_cache = self._trtllm_kv_cache_views(
+            layer.layer_id, compress_ratio
+        )
+
+        # RoPE is already applied; the unit scale makes this a plain e4m3 cast.
+        q_fp8 = q.to(torch.float8_e4m3fn).view(bs, 1, num_heads, 512)
+
+        bmm1_scale, bmm2_scale = self._get_trtllm_bmm_scales(layer)
+
+        seq_lens = core_attn_metadata.seq_lens_casual
+        if seq_lens.shape[0] != bs:
+            assert seq_lens.shape[0] > bs, f"{seq_lens.shape=} {bs=}"
+            seq_lens = seq_lens[:bs]
+        assert attn_sink.dtype == torch.float32
+        assert self.trtllm_workspace_buffer is not None
+        _check_trtllm_query_rows(bs)
+
+        out = trtllm_batch_decode_sparse_mla_dsv4(
+            query=q_fp8,
+            swa_kv_cache=swa_kv_cache,
+            workspace_buffer=self.trtllm_workspace_buffer,
+            sparse_indices=sparse_indices,
+            compressed_kv_cache=compressed_kv_cache,
+            sparse_topk_lens=sparse_topk_lens,
+            seq_lens=seq_lens,
+            bmm1_scale=bmm1_scale,
+            bmm2_scale=bmm2_scale,
+            sinks=attn_sink,
+            kv_layout="HND",
+        )
+        if out_pad_tail is not None:
+            out_pad_tail[:bs] = out.view(bs, num_heads, 512)
+            return out_pad_tail
+        return out.view(bs, num_heads, 512)
+
+    def _forward_trtllm_prefill(
+        self,
+        *,
+        q: torch.Tensor,
+        layer: RadixAttention,
+        compress_ratio: Literal[0, 4, 128],
+        forward_batch: ForwardBatch,
+        attn_sink: torch.Tensor,
+        swa_page_indices: torch.Tensor,
+        extra_indices: Optional[torch.Tensor],
+        extra_topk_lengths: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        """Drive the decode kernel as varlen prefill with one table row per token.
+
+        ``seq_lens`` includes cached prefixes; the kernel derives causal SWA
+        validity from it. This path runs eagerly.
+        """
+
+        from flashinfer.mla import trtllm_batch_decode_sparse_mla_dsv4
+
+        assert q.ndim == 3, f"{q.shape=}"
+        num_qo_padded, num_heads, head_dim = q.shape
+        assert head_dim == 512
+
+        # Build VarSeq metadata from the same extend lengths as the sparse tables.
+        core = self.forward_metadata.core_attn_metadata
+        if core.trtllm_prefill_qmeta is None:
+            extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
+            assert extend_seq_lens_cpu is not None and len(extend_seq_lens_cpu) > 0
+            batch_size = len(extend_seq_lens_cpu)
+            cum_lens = [0] * (batch_size + 1)
+            for i, extend_len in enumerate(extend_seq_lens_cpu):
+                cum_lens[i + 1] = cum_lens[i] + int(extend_len)
+            sum_q = cum_lens[-1]
+            max_q_len = max(int(x) for x in extend_seq_lens_cpu)
+            assert 0 < sum_q <= num_qo_padded, f"{sum_q=} {num_qo_padded=}"
+            seq_lens_i32 = forward_batch.seq_lens.to(torch.int32)
+            assert seq_lens_i32.shape == (batch_size,), f"{seq_lens_i32.shape=}"
+            core.trtllm_prefill_qmeta = (
+                self._move_to_device(cum_lens),
+                max_q_len,
+                sum_q,
+                seq_lens_i32,
+            )
+        cum_seq_lens_q, max_q_len, sum_q, seq_lens = core.trtllm_prefill_qmeta
+
+        # Cache layer-invariant table parts per chunk. The VarSeq kernel reads
+        # rows to a 64-token boundary, so views need inert, tile-aligned parents.
+        sum_q_pad = (sum_q + 63) // 64 * 64
+
+        def _tile_padded_pf(fill, src=None, width=None):
+            shape = (sum_q_pad,) if width is None else (sum_q_pad, width)
+            buf = torch.full(shape, fill, **self.cuda_int32_kwargs)
+            if src is not None:
+                buf[:sum_q].copy_(src)
+            return buf[:sum_q]
+
+        swa_indices = _tile_padded_pf(-1, swa_page_indices[:sum_q], width=SWA_WINDOW)
+        assert swa_indices.shape == (sum_q, SWA_WINDOW), f"{swa_indices.shape=}"
+        if extra_indices is None:
+            sparse_indices = swa_indices
+            if core.trtllm_prefill_swa_lens is None:
+                core.trtllm_prefill_swa_lens = _tile_padded_pf(SWA_WINDOW)
+            sparse_topk_lens = core.trtllm_prefill_swa_lens
+        elif compress_ratio == 128:
+            if core.trtllm_prefill_c128 is None:
+                width = extra_indices.shape[-1]
+                assert width % 4 == 0, f"{width=}"
+                table = _tile_padded_pf(-1, width=SWA_WINDOW + width)
+                table[:, :SWA_WINDOW].copy_(swa_indices)
+                table[:, SWA_WINDOW:].copy_(extra_indices[:sum_q])
+                assert extra_topk_lengths is not None
+                lens = _tile_padded_pf(
+                    SWA_WINDOW,
+                    extra_topk_lengths[:sum_q].to(torch.int32) + SWA_WINDOW,
+                )
+                core.trtllm_prefill_c128 = (table, lens)
+            sparse_indices, sparse_topk_lens = core.trtllm_prefill_c128
+        else:
+            assert extra_topk_lengths is not None
+            width = extra_indices.shape[-1]
+            # _pad_last_dim keeps the combined c4 capacity divisible by four.
+            assert width % 4 == 0, f"{width=}"
+            if core.trtllm_prefill_c4_indices is None:
+                core.trtllm_prefill_c4_indices = _tile_padded_pf(
+                    -1, width=SWA_WINDOW + width
+                )
+                core.trtllm_prefill_c4_indices[:, :SWA_WINDOW].copy_(swa_indices)
+            sparse_indices = core.trtllm_prefill_c4_indices
+            assert sparse_indices.shape == (
+                sum_q,
+                SWA_WINDOW + width,
+            ), f"{sparse_indices.shape=} {width=}"
+            sparse_indices[:, SWA_WINDOW:].copy_(extra_indices[:sum_q])
+            # Lens include 128 SWA slots; VarSeq metadata controls validity.
+            sparse_topk_lens = _tile_padded_pf(
+                SWA_WINDOW,
+                extra_topk_lengths[:sum_q].to(torch.int32) + SWA_WINDOW,
+            )
+
+        # RoPE is already applied; the unit scale makes this a plain e4m3 cast.
+        q_fp8 = q[:sum_q].to(torch.float8_e4m3fn)
+
+        swa_kv_cache, compressed_kv_cache = self._trtllm_kv_cache_views(
+            layer.layer_id, compress_ratio
+        )
+        bmm1_scale, bmm2_scale = self._get_trtllm_bmm_scales(layer)
+        assert attn_sink.dtype == torch.float32
+        assert self.trtllm_workspace_buffer is not None
+        _check_trtllm_query_rows(sum_q)
+
+        out_padded = None
+        out_arg = None
+        if num_qo_padded != sum_q:
+            # Run only real tokens and keep discarded padding rows finite.
+            out_padded = torch.zeros(
+                (num_qo_padded, num_heads, 512),
+                dtype=torch.bfloat16,
+                device=q.device,
+            )
+            out_arg = out_padded[:sum_q]
+
+        out = trtllm_batch_decode_sparse_mla_dsv4(
+            query=q_fp8,
+            swa_kv_cache=swa_kv_cache,
+            workspace_buffer=self.trtllm_workspace_buffer,
+            sparse_indices=sparse_indices,
+            compressed_kv_cache=compressed_kv_cache,
+            sparse_topk_lens=sparse_topk_lens,
+            seq_lens=seq_lens,
+            out=out_arg,
+            bmm1_scale=bmm1_scale,
+            bmm2_scale=bmm2_scale,
+            sinks=attn_sink,
+            kv_layout="HND",
+            cum_seq_lens_q=cum_seq_lens_q,
+            max_q_len=max_q_len,
+        )
+        return out_padded if out_padded is not None else out
+
+
+class DeepseekV4TrtllmMultiStepBackend(
+    DeepseekV4MultiStepBackend, DeepseekV4TrtllmAttnBackend
+):
+    """Multi-step draft wrapper whose per-step backends are trtllm."""
+
+    def _make_step_backend(
+        self, model_runner: ModelRunner, step_id: int
+    ) -> DeepseekV4AttnBackend:
+        return DeepseekV4TrtllmAttnBackend(
+            model_runner,
+            speculative_step_id=step_id,
+            topk=self.topk,
+            speculative_num_steps=self.speculative_num_steps,
+        )
+
+
+def is_dsv4_trtllm_attn_enabled() -> bool:
+    return get_exec().kernel.dsv4_attn_backend == "trtllm"
+
+
+def create_deepseek_v4_attn_backend(
+    model_runner: ModelRunner, **kwargs
+) -> DeepseekV4AttnBackend:
+    """Construct the DSV4 backend matching --dsv4-attn-backend."""
+    cls = (
+        DeepseekV4TrtllmAttnBackend
+        if is_dsv4_trtllm_attn_enabled()
+        else DeepseekV4AttnBackend
+    )
+    return cls(model_runner, **kwargs)
+
+
+def create_deepseek_v4_multistep_backend(
+    model_runner: ModelRunner, topk: int, speculative_num_steps: int
+) -> DeepseekV4MultiStepBackend:
+    cls = (
+        DeepseekV4TrtllmMultiStepBackend
+        if is_dsv4_trtllm_attn_enabled()
+        else DeepseekV4MultiStepBackend
+    )
+    return cls(model_runner, topk=topk, speculative_num_steps=speculative_num_steps)

python/sglang/srt/layers/attention/dsv4/compressor_trtllm.py

@@ -0,0 +1,120 @@
+"""Unfused compressor store for the trtllm backend's uniform-FP8 KV pool.
+
+The FlashMLA epilogue writes a different packed layout. Keep this pipeline
+separate until the fused uniform-FP8 store in PR #32975 replaces it.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.ops.attention.dsv4 import compress_forward
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.dsv4.compressor import Compressor
+    from sglang.srt.layers.attention.dsv4.compressor_v2 import CompressorBackendMixin
+    from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+
+
+def _mask_invalid_prefill_compress_rows(
+    kv_compressed: torch.Tensor,
+    plan_raw: torch.Tensor,
+    out_loc: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Make invalid prefill-plan rows safe for BCG's static-shape store."""
+    valid = plan_raw[:, 0] != -1
+    kv_compressed = torch.where(
+        valid.unsqueeze(-1), kv_compressed, torch.zeros_like(kv_compressed)
+    )
+
+    ragged_ids = plan_raw[:, 1].to(torch.int32) & 0xFFFF
+    safe_ragged_ids = torch.where(valid, ragged_ids, torch.zeros_like(ragged_ids))
+    mapped_out_loc = out_loc[safe_ragged_ids.long()]
+    # Reserved slot 0 safely absorbs duplicate invalid writes.
+    out_loc_to_store = torch.where(
+        valid, mapped_out_loc, torch.zeros_like(mapped_out_loc)
+    )
+    return kv_compressed, out_loc_to_store
+
+
+def forward_compress_uniform_fp8(
+    backend: CompressorBackendMixin,
+    *,
+    token_to_kv_pool: DeepSeekV4TokenToKVPool,
+    kv_score_input: torch.Tensor,
+    state_pool,
+    compressor: Compressor,
+    layer_id: int,
+) -> None:
+    """Compress, normalize, apply RoPE, and store as uniform e4m3."""
+    from sglang.kernels.ops.attention.deepseek_v4_rope import (
+        fused_norm_rope_inplace_triton,
+    )
+    from sglang.srt.layers.attention.dsv4.compressor_v2 import (
+        _extract_positions_from_plan,
+        _use_online_compress,
+        is_overlap_compress,
+    )
+
+    assert not compressor.is_in_indexer
+    assert compressor.head_dim == 512, f"{compressor.head_dim=}"
+    assert not _use_online_compress(compressor.ratio), (
+        "SGLANG_OPT_USE_ONLINE_COMPRESS is not supported with the "
+        "uniform-FP8 KV layout yet."
+    )
+
+    compress_ratio = compressor.ratio
+    head_dim = compressor.head_dim
+
+    plan = backend._get_paged_compress_metadata(compress_ratio)
+    out_loc = backend._get_out_loc(compress_ratio)
+
+    coff = 2 if is_overlap_compress(compress_ratio) else 1
+    kv_score_buffer = state_pool.kv_score_buffer.kv_score.view(
+        -1, compress_ratio, 2 * head_dim * coff
+    )
+
+    kv_compressed = compress_forward(
+        kv_score_buffer=kv_score_buffer,
+        kv_score_input=kv_score_input,
+        ape=compressor.ape.view(-1, head_dim),
+        plan=plan,
+        compress_ratio=compress_ratio,
+        head_dim=head_dim,
+        is_online=False,
+    )
+    if kv_compressed.shape[0] == 0:
+        return
+
+    plan_raw = plan[1].view(torch.int32)
+    if plan.is_decode:
+        # Zero out non-boundary tokens to prevent corrupting kvcache loc 0.
+        seq_lens_plan = plan_raw[:, 0].to(torch.int32)
+        is_boundary = (seq_lens_plan % compress_ratio == 0).unsqueeze(-1)
+        kv_compressed = torch.where(
+            is_boundary, kv_compressed, torch.zeros_like(kv_compressed)
+        )
+        out_loc_to_store = out_loc
+    else:
+        kv_compressed, out_loc_to_store = _mask_invalid_prefill_compress_rows(
+            kv_compressed,
+            plan_raw,
+            out_loc,
+        )
+
+    positions = _extract_positions_from_plan(plan, compress_ratio).clamp(min=0)
+    fused_norm_rope_inplace_triton(
+        kv_compressed,
+        compressor.norm.weight,
+        compressor.norm.variance_epsilon,
+        compressor.freqs_cis,
+        positions=positions,
+    )
+
+    token_to_kv_pool.set_extra_key_buffer_fused(
+        layer_id=layer_id,
+        loc=out_loc_to_store,
+        cache_k=kv_compressed,
+    )

python/sglang/srt/layers/attention/dsv4/compressor_v2.py

@@ -241,6 +241,22 @@ class CompressorBackendMixin:
             is_unified_kv_triton,
         )
 
+        if token_to_kv_pool.uniform_fp8 and not compressor.is_in_indexer:
+            # The fused epilogue writes only the packed FlashMLA layout.
+            from sglang.srt.layers.attention.dsv4.compressor_trtllm import (
+                forward_compress_uniform_fp8,
+            )
+
+            forward_compress_uniform_fp8(
+                self,
+                token_to_kv_pool=token_to_kv_pool,
+                kv_score_input=kv_score_input,
+                state_pool=state_pool,
+                compressor=compressor,
+                layer_id=layer_id,
+            )
+            return
+
         out_loc = self._get_out_loc(compressor.ratio)
         use_fp4_indexer = (
             compressor.is_in_indexer and self.enable_deepseek_v4_fp4_indexer

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -187,6 +187,63 @@ class DeepSeekV4SingleKVPool(KVCache):
         raise NotImplementedError("Use get_key_buffer instead.")
 
 
+class DeepSeekV4UniformFP8KVPool(DeepSeekV4SingleKVPool):
+    """Uniform 512-dim FP8 (e4m3) variant of the DSv4 single-KV pool.
+
+    Each token is 448 NoPE + 64 RoPE contiguous e4m3 values without in-cache
+    scales or per-page padding. The backend supplies the dequant scale.
+    """
+
+    def get_bytes_per_token(self) -> int:
+        return self.qk_nope_head_dim + self.qk_rope_head_dim
+
+    def create_buffer(self, *, num_pages: int):
+        bytes_per_token = self.get_bytes_per_token()
+        assert bytes_per_token == 512, (
+            "DSV4 uniform-FP8 KV layout: qk_nope_head_dim (448) + "
+            "qk_rope_head_dim (64), all e4m3 = 512 bytes/token"
+        )
+        self.kv_cache_total_dim = bytes_per_token
+        self.bytes_per_page_padded = self.page_size * bytes_per_token
+
+        return torch.zeros(
+            num_pages,
+            self.page_size * bytes_per_token,
+            dtype=torch.float8_e4m3fn,
+            device=self.device,
+        )
+
+    def get_key_buffer(self, layer_id: int):
+        return self.kv_buffer[layer_id]
+
+    def set_key_buffer(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        cache_nope_fp8_rope_bf16_pack: NopeFp8RopeBf16Pack,
+    ):
+        raise NotImplementedError(
+            "The packed NopeFp8RopeBf16Pack store does not apply to the "
+            "uniform-FP8 pool; use set_key_buffer_fused."
+        )
+
+    def set_key_buffer_fused(
+        self,
+        layer_id: int,
+        loc: torch.Tensor,
+        cache_k: torch.Tensor,
+    ) -> None:
+        """Store normed/roped rows as e4m3 with the backend's fixed unit scale.
+
+        uint8 views work around index_put not supporting FP8 dtypes.
+        """
+
+        assert cache_k.dim() == 2 and cache_k.shape[1] == self.kv_cache_total_dim
+        self.kv_buffer[layer_id].view(torch.uint8).view(-1, self.kv_cache_total_dim)[
+            loc.long()
+        ] = cache_k.to(torch.float8_e4m3fn).view(torch.uint8)
+
+
 class HiSparseC4DevicePool(DeepSeekV4SingleKVPool):
     def __init__(
         self,
@@ -578,6 +635,10 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
 
         # Resolve the unified-kv gate before any sizing so the two cannot drift.
         self._unified_kv = is_unified_kv_triton()
+        # Uniform 512-dim e4m3 layout for the trtllm attention backend
+        self.uniform_fp8 = (
+            not self._unified_kv
+        ) and get_exec().kernel.dsv4_attn_backend == "trtllm"
         c4_ring_size = self.get_ring_size(4)
         if self._unified_kv:
             # Unified C4 state is request-addressed: one ring per req slot,
@@ -665,6 +726,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
             self.swa_req_ring_size = self.unified_swa_ring_size
         else:
             self.unified_kv_pool = None
+            kv_pool_cls: type = DeepSeekV4SingleKVPool
+            if self.uniform_fp8:
+                assert dtype == torch.float8_e4m3fn, (
+                    "--dsv4-attn-backend trtllm requires "
+                    f"kv_cache_dtype=fp8_e4m3, got {dtype}"
+                )
+                kv_pool_cls = DeepSeekV4UniformFP8KVPool
             self.swa_kv_pool = self._make_kv_pool(
                 size=swa_size,
                 page_size=swa_page_size,
@@ -673,10 +741,14 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
                 device=device,
                 enable_memory_saver=enable_memory_saver,
                 global_page_size=swa_page_size,
+                cls=kv_pool_cls,
             )
 
-            c4_kv_pool_type = DeepSeekV4SingleKVPool
+            c4_kv_pool_type = kv_pool_cls
             if enable_hisparse:
+                assert not self.uniform_fp8, (
+                    "enable_hisparse is not supported with --dsv4-attn-backend trtllm."
+                )
                 c4_kv_pool_type = HiSparseC4DevicePool
             self.c4_kv_pool = self._make_kv_pool(
                 size=c4_size,
@@ -697,6 +769,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
                 device=device,
                 enable_memory_saver=enable_memory_saver,
                 global_page_size=page_size,
+                cls=kv_pool_cls,
             )
 
         indexer_size = self.c4_logical_size
@@ -1286,6 +1359,26 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         freqs_cis: torch.Tensor,
         positions: torch.Tensor,
     ) -> None:
+        if self.uniform_fp8:
+            # Uniform-FP8 (trtllm-gen) layout: norm + RoPE with the existing
+            # Triton kernel (in-place on kv; safe -- kv is not read again),
+            # then a plain e4m3 cast + scatter in the pool setter (per-tensor
+            # scale 1.0). Fusing the store is deferred to the perf phase.
+            from sglang.kernels.ops.attention.deepseek_v4_rope import (
+                fused_norm_rope_inplace_triton,
+            )
+
+            fused_norm_rope_inplace_triton(
+                kv,
+                kv_weight,
+                eps,
+                freqs_cis,
+                positions=positions,
+            )
+            self.swa_kv_pool.set_key_buffer_fused(
+                self._swa_local_layer_id(layer_id), swa_loc, kv
+            )
+            return
         fused_k_norm_rope_flashmla(
             kv=kv,
             kv_weight=kv_weight,

python/sglang/srt/speculative/draft_utils.py

@@ -433,8 +433,8 @@ class DraftBackendFactory:
                 DeepseekV4MultiStepBackend,
             )
         else:
-            from sglang.srt.layers.attention.deepseek_v4_backend import (
+            from sglang.srt.layers.attention.deepseek_v4_trtllm_backend import (
-                DeepseekV4MultiStepBackend,
+                create_deepseek_v4_multistep_backend as DeepseekV4MultiStepBackend,
             )
 
         return (
@@ -573,11 +573,13 @@ class DraftBackendFactory:
                 "dsv4",
                 DeepseekV4HipRadixBackend(self.draft_model_runner, skip_prefill=False),
             )
-        from sglang.srt.layers.attention.deepseek_v4_backend import (
+        from sglang.srt.layers.attention.deepseek_v4_trtllm_backend import (
-            DeepseekV4AttnBackend,
+            create_deepseek_v4_attn_backend,
         )
 
         return (
             "dsv4",
-            DeepseekV4AttnBackend(self.draft_model_runner, skip_prefill=False),
+            create_deepseek_v4_attn_backend(
+                self.draft_model_runner, skip_prefill=False
+            ),
         )

test/registered/attention/unittests/dsv4/test_deepseek_v4.py

@@ -604,6 +604,34 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase):
             )
         )
 
+    def test_trtllm_semaphore_capacity_covers_configured_query_rows(self):
+        from sglang.srt.layers.attention import deepseek_v4_trtllm_backend as trtllm
+
+        schedule = SimpleNamespace(
+            max_prefill_tokens=16384,
+            chunked_prefill_size=4096,
+            max_running_requests=256,
+        )
+        spec = SimpleNamespace(
+            speculative_algorithm="EAGLE", speculative_num_draft_tokens=4
+        )
+        model_runner = SimpleNamespace()
+        with (
+            mock.patch.object(trtllm, "get_schedule", return_value=schedule),
+            mock.patch.object(trtllm, "get_spec", return_value=spec),
+            mock.patch.object(trtllm, "max_prefill_buffer_tokens", return_value=4096),
+        ):
+            # Prefill chunk / max_prefill_tokens dominates.
+            self.assertEqual(trtllm._trtllm_query_row_capacity(model_runner), 16384)
+            # Decode rows = requests x draft tokens dominate.
+            schedule.max_running_requests = 8192
+            self.assertEqual(trtllm._trtllm_query_row_capacity(model_runner), 32768)
+
+        with mock.patch.object(trtllm, "_trtllm_semaphore_rows", 64):
+            trtllm._check_trtllm_query_rows(64)
+            with self.assertRaisesRegex(RuntimeError, "exceeds the persistent"):
+                trtllm._check_trtllm_query_rows(65)
+
     def test_sparse_prefill_workspace_reuses_and_grows(self):
         from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
             SparsePrefillWorkspace,

test/registered/e2e/dsv4/test_dsv4_fp8_trtllm_backend.py

@@ -0,0 +1,214 @@
+"""SM100/SM103 coverage for DSV4's uniform-FP8 trtllm backend.
+
+Covers decode correctness, GSM8K accuracy, varlen and cached-prefix prefill,
+chunking, and decode CUDA-graph replay. Long outputs use sanity checks because
+the FlashMLA and uniform-FP8 cache formats need not be bit-reproducible.
+"""
+
+import concurrent.futures
+import unittest
+from types import SimpleNamespace
+
+import requests
+import torch
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+    try_cached_model,
+)
+
+register_cuda_ci(est_time=900, stage="base-c", runner_config="4-gpu-b200")
+
+DSV4_FLASH_MODEL_PATH = try_cached_model("deepseek-ai/DeepSeek-V4-Flash")
+SERVER_LAUNCH_TIMEOUT = 3600
+DSV4_BASE_ENV = {
+    "SGLANG_JIT_DEEPGEMM_FAST_WARMUP": "1",
+}
+
+SERVER_ARGS = [
+    "--trust-remote-code",
+    "--dsv4-attn-backend",
+    "trtllm",
+    "--tp",
+    "4",
+    "--max-running-requests",
+    "32",
+    "--mem-fraction-static",
+    "0.85",
+    "--chunked-prefill-size",
+    "4096",
+    # V4-Flash ships MXFP4 routed experts, and the auto-selected Triton MoE runner
+    # cannot consume the packed layout. Matches the B200 Flash cookbook recipe.
+    "--moe-runner-backend",
+    "flashinfer_mxfp4",
+    "--disable-flashinfer-autotune",
+]
+
+# Mixed lengths cover c4/c128 selection and VarSeq packing; the longest prompt
+# exceeds the 4096-token prefill chunk.
+_FILLER_SENTENCES = [
+    "The expedition recorded water temperature, salinity, and current speed "
+    "at every station along the transect. ",
+    "Archival records from the observatory describe decades of nightly "
+    "measurements taken with remarkable consistency. ",
+    "Each greenhouse module recycles condensate through a gravel bed before "
+    "returning it to the irrigation loop. ",
+    "The survey team catalogued the masonry of the aqueduct arch by arch, "
+    "noting repairs from three distinct centuries. ",
+]
+_LONG_PROMPT_QUESTION = (
+    "\n\nIn one short sentence, what kind of activity do the paragraphs above describe?"
+)
+
+
+def _make_long_prompt(idx: int, target_chars: int) -> str:
+    sentence = _FILLER_SENTENCES[idx % len(_FILLER_SENTENCES)]
+    body = ""
+    n = 0
+    while len(body) < target_chars:
+        body += f"[Entry {idx}-{n}] " + sentence
+        n += 1
+    return body + _LONG_PROMPT_QUESTION
+
+
+# Roughly 2.5k, 4.5k, and 7k tokens.
+LONG_PROMPTS = [
+    _make_long_prompt(0, 10_000),
+    _make_long_prompt(1, 18_000),
+    _make_long_prompt(2, 28_000),
+]
+LONG_MAX_NEW_TOKENS = 32
+MIN_PRINTABLE_ASCII_RATIO = 0.85
+
+GSM8K_NUM_EXAMPLES = 200
+GSM8K_MIN_SCORE = 0.90
+
+_REQUEST_TIMEOUT = 600
+
+
+def _is_sm100() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    return torch.cuda.get_device_capability() in ((10, 0), (10, 3))
+
+
+def _greedy_generate(base_url: str, prompt: str, max_new_tokens: int) -> str:
+    resp = requests.post(
+        base_url + "/generate",
+        json={
+            "text": prompt,
+            "sampling_params": {
+                "temperature": 0.0,
+                "max_new_tokens": max_new_tokens,
+            },
+        },
+        timeout=_REQUEST_TIMEOUT,
+    )
+    resp.raise_for_status()
+    return resp.json()["text"]
+
+
+def _printable_ascii_ratio(text: str) -> float:
+    if not text:
+        return 0.0
+    return sum(32 <= ord(c) < 127 or c in "\n\t" for c in text) / len(text)
+
+
+class TestDSV4Fp8TrtllmBackend(BasicDecodeCorrectnessMixin, CustomTestCase):
+    """TP4 DSv4-Flash-FP8 with --dsv4-attn-backend trtllm."""
+
+    @classmethod
+    def setUpClass(cls):
+        if not _is_sm100():
+            raise unittest.SkipTest(
+                "DSv4 trtllm uniform-FP8 attention requires SM100/SM103 (Blackwell)"
+            )
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            DSV4_FLASH_MODEL_PATH,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=SERVER_ARGS,
+            env=dict(DSV4_BASE_ENV),
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process is not None:
+            kill_process_tree(cls.process.pid)
+
+    def _assert_sane(self, out: str, what: str) -> None:
+        self.assertGreater(len(out.strip()), 0, f"{what}: empty output")
+        ratio = _printable_ascii_ratio(out)
+        self.assertGreater(
+            ratio,
+            MIN_PRINTABLE_ASCII_RATIO,
+            f"{what}: output looks like gibberish (ascii ratio={ratio:.2f}): {out!r}",
+        )
+
+    def test_long_prompt_varlen_prefill(self):
+        """Exercise mixed-length VarSeq and cached-prefix chunked prefill.
+
+        Sanity checks avoid flaky exact matches from split-KV reduction order.
+        """
+
+        with concurrent.futures.ThreadPoolExecutor(len(LONG_PROMPTS)) as pool:
+            outs = list(
+                pool.map(
+                    lambda p: _greedy_generate(self.base_url, p, LONG_MAX_NEW_TOKENS),
+                    LONG_PROMPTS,
+                )
+            )
+        for i, out in enumerate(outs):
+            print(f"[long-prefill] prompt_chars={len(LONG_PROMPTS[i])} out={out!r}")
+            self._assert_sane(out, f"concurrent long prompt {i}")
+
+        cached = _greedy_generate(self.base_url, LONG_PROMPTS[-1], LONG_MAX_NEW_TOKENS)
+        print(f"[long-prefill] cached-prefix rerun out={cached!r}")
+        self._assert_sane(cached, "cached-prefix extend")
+
+    def test_gsm8k_sanity(self):
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=DSV4_FLASH_MODEL_PATH,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=GSM8K_NUM_EXAMPLES,
+            num_threads=64,
+        )
+        metrics = run_eval(args)
+        print(f"GSM8K sanity on trtllm decode: {metrics=}")
+        self.assertGreater(metrics["score"], GSM8K_MIN_SCORE)
+
+    def test_cuda_graph_capture_replay_smoke(self):
+        """Replay several decode graph buckets and recheck a greedy anchor."""
+        anchor_prompt = "Q: What is the capital of France?\nA:"
+        anchor_out = _greedy_generate(self.base_url, anchor_prompt, 32)
+
+        for concurrency in (2, 4, 8, 16):
+            prompts = [f"Count from {i} to {i + 5}: " for i in range(concurrency)]
+            with concurrent.futures.ThreadPoolExecutor(concurrency) as pool:
+                outs = list(
+                    pool.map(lambda p: _greedy_generate(self.base_url, p, 32), prompts)
+                )
+            self.assertEqual(len(outs), concurrency)
+            for out in outs:
+                self.assertGreater(len(out), 0)
+
+        anchor_out_replayed = _greedy_generate(self.base_url, anchor_prompt, 32)
+        self.assertEqual(
+            anchor_out,
+            anchor_out_replayed,
+            "greedy output changed after batched decode-graph replays",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=3)

test/registered/e2e/models/test_deepseek_v4_flash_fp4_b200_trtllm.py

@@ -0,0 +1,223 @@
+"""B200 per-commit CI: DeepSeek-V4-Flash FP4 with the trtllm attention backend.
+
+Mirrors the four FlashMLA recipes with a uniform-FP8 KV pool and trtllm-gen
+sparse MLA for decode and prefill.
+"""
+
+import unittest
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.kits.spec_decoding_kit import SpecDecodingMixin
+from sglang.test.test_utils import (
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+    try_cached_model,
+)
+
+register_cuda_ci(est_time=700, stage="base-c", runner_config="4-gpu-b200")
+
+MODEL = "deepseek-ai/DeepSeek-V4-Flash"
+SERVER_LAUNCH_TIMEOUT = 3600
+DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
+
+_DEEPEP_ENV = {
+    "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
+}
+
+
+class TestDSV4FlashFP4B200Trtllm(
+    SpecDecodingMixin,
+    BasicDecodeCorrectnessMixin,
+    GSM8KMixin,
+    CustomTestCase,
+):
+    """LowLatency recipe: TP=4, FP4 (mxfp4), EAGLE spec decoding."""
+
+    gsm8k_accuracy_thres = 0.93
+    accept_length_thres = 2.8
+    bs_1_speed_thres = 220
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = try_cached_model(MODEL)
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=[
+                "--trust-remote-code",
+                "--dsv4-attn-backend",
+                "trtllm",
+                "--tp",
+                "4",
+                "--moe-runner-backend",
+                "flashinfer_mxfp4",
+                "--speculative-algorithm",
+                "EAGLE",
+                "--speculative-num-steps",
+                "3",
+                "--speculative-eagle-topk",
+                "1",
+                "--speculative-num-draft-tokens",
+                "4",
+                "--chunked-prefill-size",
+                "4096",
+                "--disable-flashinfer-autotune",
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+
+class TestDSV4FlashFP4B200BalancedTrtllm(
+    SpecDecodingMixin,
+    BasicDecodeCorrectnessMixin,
+    GSM8KMixin,
+    CustomTestCase,
+):
+    """Balanced recipe: TP=4, DP=4, DeepEP, EAGLE (1-step spec)."""
+
+    gsm8k_accuracy_thres = 0.93
+    accept_length_thres = 1.8
+    bs_1_speed_thres = 100
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = try_cached_model(MODEL)
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=[
+                "--trust-remote-code",
+                "--dsv4-attn-backend",
+                "trtllm",
+                "--tp",
+                "4",
+                "--dp",
+                "4",
+                "--enable-dp-attention",
+                "--moe-a2a-backend",
+                "deepep",
+                "--speculative-algorithm",
+                "EAGLE",
+                "--speculative-num-steps",
+                "1",
+                "--speculative-eagle-topk",
+                "1",
+                "--speculative-num-draft-tokens",
+                "2",
+                "--deepep-config",
+                DEEPEP_CONFIG,
+            ],
+            env=_DEEPEP_ENV,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+
+class TestDSV4FlashFP4NonMTPB200Trtllm(
+    BasicDecodeCorrectnessMixin, GSM8KMixin, CustomTestCase
+):
+    """Non-MTP recipe: TP=4, DP=4, DeepEP, no speculative decoding."""
+
+    gsm8k_accuracy_thres = 0.93
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = try_cached_model(MODEL)
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=[
+                "--trust-remote-code",
+                "--dsv4-attn-backend",
+                "trtllm",
+                "--tp",
+                "4",
+                "--dp",
+                "4",
+                "--enable-dp-attention",
+                "--moe-a2a-backend",
+                "deepep",
+                "--deepep-config",
+                DEEPEP_CONFIG,
+            ],
+            env=_DEEPEP_ENV,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+
+class TestDSV4FlashFP4BreakableCudaGraphB200Trtllm(
+    BasicDecodeCorrectnessMixin, GSM8KMixin, CustomTestCase
+):
+    """BCG recipe: TP=4, DP=4, DeepEP, DP attention, mixed chunk."""
+
+    gsm8k_accuracy_thres = 0.93
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = try_cached_model(MODEL)
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=SERVER_LAUNCH_TIMEOUT,
+            other_args=[
+                "--trust-remote-code",
+                "--dsv4-attn-backend",
+                "trtllm",
+                "--tp",
+                "4",
+                "--dp",
+                "4",
+                "--enable-dp-attention",
+                "--enable-mixed-chunk",
+                "--cuda-graph-backend-prefill",
+                "breakable",
+                "--moe-a2a-backend",
+                "deepep",
+                "--deepep-config",
+                DEEPEP_CONFIG,
+                "--chunked-prefill-size",
+                "4096",
+                "--cuda-graph-max-bs-prefill",
+                "1024",
+                "--mem-fraction-static",
+                "0.80",
+                "--cuda-graph-max-bs-decode",
+                "16",
+                "--max-running-requests",
+                "128",
+                "--watchdog-timeout",
+                "900",
+            ],
+            env=_DEEPEP_ENV,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/test_model_overrides.py

@@ -95,6 +95,7 @@ class TestModelOverridableWhitelist(CustomTestCase):
                     "kv_cache_dtype",
                     "dsa_prefill_backend",
                     "dsa_decode_backend",
+                    "dsv4_attn_backend",
                     "dsa_topk_backend",
                     "prefill_attention_backend",
                     "decode_attention_backend",