[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>
This commit is contained in:
akhilg-nv
co-authored by
Yangmin Li
Po-Han Huang
Yangmin Li
Po-Han Huang
parent
0084030179
commit
880d6fa64d
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user