From 880d6fa64d187ae1d14b58101e2a776853492f26 Mon Sep 17 00:00:00 2001 From: akhilg-nv <165961486+akhilg-nv@users.noreply.github.com> Date: Wed, 9 Sep 2026 17:40:07 -0700 Subject: [PATCH] [DSv4] Integrate TRT-LLM DSv4 Attention for SM100/103 (#30805) Co-authored-by: Yangmin Li Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com> --- .../sglang/srt/arg_groups/deepseek_v4_hook.py | 41 ++ python/sglang/srt/arg_groups/fields/exec_.py | 13 + .../layers/attention/attention_registry.py | 11 +- .../layers/attention/deepseek_v4_backend.py | 120 +++- .../attention/deepseek_v4_trtllm_backend.py | 532 ++++++++++++++++++ .../attention/dsv4/compressor_trtllm.py | 120 ++++ .../layers/attention/dsv4/compressor_v2.py | 16 + .../srt/mem_cache/deepseek_v4_memory_pool.py | 95 +++- python/sglang/srt/speculative/draft_utils.py | 12 +- .../unittests/dsv4/test_deepseek_v4.py | 28 + .../e2e/dsv4/test_dsv4_fp8_trtllm_backend.py | 214 +++++++ .../test_deepseek_v4_flash_fp4_b200_trtllm.py | 223 ++++++++ test/registered/unit/test_model_overrides.py | 1 + 13 files changed, 1408 insertions(+), 18 deletions(-) create mode 100644 python/sglang/srt/layers/attention/deepseek_v4_trtllm_backend.py create mode 100644 python/sglang/srt/layers/attention/dsv4/compressor_trtllm.py create mode 100644 test/registered/e2e/dsv4/test_dsv4_fp8_trtllm_backend.py create mode 100644 test/registered/e2e/models/test_deepseek_v4_flash_fp4_b200_trtllm.py diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 731f8f5a1..ce08fd811 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/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, diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index 289a08ef1..2c05b9957 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/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( diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 4466b55c3..30ff4f228 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/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 ( - DeepseekV4AttnBackend, + from sglang.srt.layers.attention.deepseek_v4_trtllm_backend import ( + create_deepseek_v4_attn_backend, ) - logger.info("Using DeepseekV4AttnBackend for dsv4 attention backend (CUDA).") - return DeepseekV4AttnBackend(runner) + backend = create_deepseek_v4_attn_backend(runner) + logger.info( + f"Using {type(backend).__name__} for dsv4 attention backend (CUDA)." + ) + return backend @register_attention_backend("triton") diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 5827504d4..799834048 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/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( - DeepseekV4AttnBackend( - model_runner, - speculative_step_id=i, - topk=self.topk, - speculative_num_steps=self.speculative_num_steps, - ) - ) + self.attn_backends.append(self._make_step_backend(model_runner, i)) + + def _make_step_backend( + self, model_runner: ModelRunner, step_id: int + ) -> DeepseekV4AttnBackend: + return DeepseekV4AttnBackend( + model_runner, + 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: diff --git a/python/sglang/srt/layers/attention/deepseek_v4_trtllm_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_trtllm_backend.py new file mode 100644 index 000000000..634c6a63e --- /dev/null +++ b/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) diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_trtllm.py b/python/sglang/srt/layers/attention/dsv4/compressor_trtllm.py new file mode 100644 index 000000000..ecdcd2c46 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index c35154131..207fdb7e5 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 455847942..ddc7c7eca 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/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, diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py index f93f2f648..365c11b81 100644 --- a/python/sglang/srt/speculative/draft_utils.py +++ b/python/sglang/srt/speculative/draft_utils.py @@ -433,8 +433,8 @@ class DraftBackendFactory: DeepseekV4MultiStepBackend, ) else: - from sglang.srt.layers.attention.deepseek_v4_backend import ( - DeepseekV4MultiStepBackend, + from sglang.srt.layers.attention.deepseek_v4_trtllm_backend import ( + 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 ( - DeepseekV4AttnBackend, + from sglang.srt.layers.attention.deepseek_v4_trtllm_backend import ( + 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 + ), ) diff --git a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py index 0f2fd7503..20d1a74fc 100644 --- a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py +++ b/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, diff --git a/test/registered/e2e/dsv4/test_dsv4_fp8_trtllm_backend.py b/test/registered/e2e/dsv4/test_dsv4_fp8_trtllm_backend.py new file mode 100644 index 000000000..871b95b0b --- /dev/null +++ b/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) diff --git a/test/registered/e2e/models/test_deepseek_v4_flash_fp4_b200_trtllm.py b/test/registered/e2e/models/test_deepseek_v4_flash_fp4_b200_trtllm.py new file mode 100644 index 000000000..3f9ceeaf3 --- /dev/null +++ b/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() diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 722be661b..51c8f1e13 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/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",