[DeepSeek-V4] Add Q8KV8 sparse MLA prefill runtime backend (#32327)
Co-authored-by: Ho-Ren (Jack) Chuang <horenchuang@bytedance.com> Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
shiyang814-cpu
co-authored by
Ho-Ren Chuang
Xiaoyu Zhang
Ho-Ren Chuang
Xiaoyu Zhang
parent
a49560ce50
commit
9db4ba8da1
@@ -85,6 +85,140 @@ def dequantize_k_cache_paged(
|
||||
return out
|
||||
|
||||
|
||||
def gather_dequant_requant_fp8_paged(
|
||||
quant_k_cache: torch.Tensor,
|
||||
page_table_1_flattened: torch.Tensor,
|
||||
page_size: int,
|
||||
extra_rows: int = 0,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
"""Gather DeepSeek-V4 paged KV cache into a flat FP8 workspace.
|
||||
|
||||
This is the Q8KV8 sparse-prefill adapter for the DeepSeek-V4 packed layout.
|
||||
It gathers token IDs from the existing paged cache, dequantizes the 448-dim
|
||||
nope region with its UE8M0 per-64 scales, casts the 64-dim BF16 rope tail to
|
||||
FP8, and writes the result as ``(num_tokens + extra_rows, 1, 512)`` FP8.
|
||||
|
||||
``extra_rows`` appends zero rows for kernels that map masked sparse indices
|
||||
to a valid zero landing pad.
|
||||
"""
|
||||
assert quant_k_cache.is_contiguous()
|
||||
assert page_table_1_flattened.dtype in (torch.int32, torch.int64)
|
||||
assert extra_rows >= 0
|
||||
|
||||
quant_k_cache_u8 = quant_k_cache.view(torch.uint8)
|
||||
num_tokens = page_table_1_flattened.shape[0]
|
||||
total_rows = num_tokens + extra_rows
|
||||
bytes_per_page = quant_k_cache_u8.shape[-1]
|
||||
s_offset_bytes = page_size * NOPE_ROPE_BYTES
|
||||
|
||||
buf_fp8 = quant_k_cache_u8.view(fp8_dtype).reshape(-1)
|
||||
buf_bf16 = quant_k_cache_u8.view(torch.bfloat16).reshape(-1)
|
||||
buf_uint8 = quant_k_cache_u8.reshape(-1)
|
||||
|
||||
if out is None:
|
||||
out = torch.zeros(
|
||||
(total_rows, 1, DIM_NOPE + DIM_ROPE),
|
||||
dtype=fp8_dtype,
|
||||
device=quant_k_cache.device,
|
||||
)
|
||||
else:
|
||||
assert out.shape == (total_rows, 1, DIM_NOPE + DIM_ROPE)
|
||||
assert out.dtype == fp8_dtype
|
||||
if extra_rows:
|
||||
out[num_tokens:].zero_()
|
||||
|
||||
if num_tokens == 0:
|
||||
return out
|
||||
|
||||
_gather_dequant_requant_fp8_paged_kernel[(num_tokens,)](
|
||||
out,
|
||||
buf_fp8,
|
||||
buf_bf16,
|
||||
buf_uint8,
|
||||
page_table_1_flattened,
|
||||
out.stride(0),
|
||||
BYTES_PER_PAGE=bytes_per_page,
|
||||
PAGE_SIZE=page_size,
|
||||
DIM_NOPE=DIM_NOPE,
|
||||
DIM_ROPE=DIM_ROPE,
|
||||
TILE_SIZE=TILE_SIZE,
|
||||
NUM_SCALE_TILES=NUM_SCALE_TILES,
|
||||
NOPE_ROPE_BYTES=NOPE_ROPE_BYTES,
|
||||
PADDED_SCALE_PER_TOKEN=PADDED_SCALE_PER_TOKEN,
|
||||
S_OFFSET_BYTES=s_offset_bytes,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def q8kv8_padded_num_heads(num_heads: int) -> int:
|
||||
"""Return a Q-head count supported by the SM90 Q8KV8 kernel."""
|
||||
if num_heads <= 0:
|
||||
raise ValueError(f"num_heads must be positive, got {num_heads}")
|
||||
if num_heads <= 64:
|
||||
return 64
|
||||
if num_heads <= 128:
|
||||
return 128
|
||||
raise ValueError(
|
||||
"DeepSeek-V4 Q8KV8 sparse prefill supports at most 128 local "
|
||||
f"query heads, got {num_heads}"
|
||||
)
|
||||
|
||||
|
||||
def cast_q_fp8_for_q8kv8_prefill(
|
||||
q: torch.Tensor,
|
||||
padded_num_heads: Optional[int] = None,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Cast DeepSeek-V4 sparse-prefill Q to the Q8KV8 kernel format.
|
||||
|
||||
The incoming Q is the model-produced BF16/FP16 tensor already shaped as
|
||||
``(num_tokens, num_heads, 512)`` after removing the singleton MQA axis.
|
||||
|
||||
The SM90 kernel processes query heads in 64-head blocks. Tensor parallelism
|
||||
commonly leaves fewer than 64 local heads, so the active heads are copied
|
||||
into a zero-padded 64/128-head FP8 tensor.
|
||||
"""
|
||||
assert q.ndim == 3
|
||||
assert q.shape[-1] == DIM_NOPE + DIM_ROPE
|
||||
|
||||
num_tokens, num_heads, head_dim = q.shape
|
||||
if padded_num_heads is None:
|
||||
padded_num_heads = q8kv8_padded_num_heads(num_heads)
|
||||
|
||||
if padded_num_heads not in (64, 128) or padded_num_heads < num_heads:
|
||||
raise ValueError(
|
||||
f"invalid padded_num_heads={padded_num_heads} for num_heads={num_heads}"
|
||||
)
|
||||
|
||||
expected_shape = (num_tokens, padded_num_heads, head_dim)
|
||||
|
||||
if out is None:
|
||||
q_fp8 = torch.zeros(
|
||||
expected_shape,
|
||||
dtype=fp8_dtype,
|
||||
device=q.device,
|
||||
)
|
||||
else:
|
||||
if (
|
||||
out.shape != expected_shape
|
||||
or out.dtype != fp8_dtype
|
||||
or out.device != q.device
|
||||
):
|
||||
raise ValueError(
|
||||
"Q8KV8 Q output must have shape/dtype/device "
|
||||
f"{expected_shape}/{fp8_dtype}/{q.device}, got "
|
||||
f"{tuple(out.shape)}/{out.dtype}/{out.device}"
|
||||
)
|
||||
q_fp8 = out
|
||||
if padded_num_heads > num_heads:
|
||||
q_fp8[:, num_heads:].zero_()
|
||||
|
||||
q_fp8[:, :num_heads].copy_(q)
|
||||
q_scale = torch.ones((), dtype=torch.float32, device=q.device)
|
||||
return q_fp8, q_scale
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _dequantize_k_cache_paged_kernel(
|
||||
output_ptr,
|
||||
@@ -136,6 +270,58 @@ def _dequantize_k_cache_paged_kernel(
|
||||
tl.store(output_ptr + out_row_base + DIM_NOPE + rope_offs, rope_data)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _gather_dequant_requant_fp8_paged_kernel(
|
||||
output_ptr,
|
||||
buf_fp8_ptr,
|
||||
buf_bf16_ptr,
|
||||
buf_uint8_ptr,
|
||||
page_table_ptr,
|
||||
output_stride_0,
|
||||
BYTES_PER_PAGE: tl.constexpr,
|
||||
PAGE_SIZE: tl.constexpr,
|
||||
DIM_NOPE: tl.constexpr,
|
||||
DIM_ROPE: tl.constexpr,
|
||||
TILE_SIZE: tl.constexpr,
|
||||
NUM_SCALE_TILES: tl.constexpr,
|
||||
NOPE_ROPE_BYTES: tl.constexpr,
|
||||
PADDED_SCALE_PER_TOKEN: tl.constexpr,
|
||||
S_OFFSET_BYTES: tl.constexpr,
|
||||
):
|
||||
token_id = tl.program_id(0)
|
||||
loc = tl.load(page_table_ptr + token_id).to(tl.int64)
|
||||
page_idx = loc // PAGE_SIZE
|
||||
in_page = loc % PAGE_SIZE
|
||||
page_byte_base = page_idx * BYTES_PER_PAGE
|
||||
token_data_base = page_byte_base + in_page * NOPE_ROPE_BYTES
|
||||
token_scale_base = (
|
||||
page_byte_base + S_OFFSET_BYTES + in_page * PADDED_SCALE_PER_TOKEN
|
||||
)
|
||||
out_row_base = token_id * output_stride_0
|
||||
|
||||
nope_offs = tl.arange(0, TILE_SIZE)
|
||||
for tile_id in tl.static_range(NUM_SCALE_TILES):
|
||||
fp8_off = token_data_base + tile_id * TILE_SIZE + nope_offs
|
||||
fp8_vals = tl.load(buf_fp8_ptr + fp8_off).to(tl.float32)
|
||||
|
||||
scale_u8 = tl.load(buf_uint8_ptr + token_scale_base + tile_id).to(tl.int32)
|
||||
scale_pow2 = tl.exp2((scale_u8 - 127).to(tl.float32))
|
||||
|
||||
out_off = out_row_base + tile_id * TILE_SIZE + nope_offs
|
||||
tl.store(
|
||||
output_ptr + out_off,
|
||||
(fp8_vals * scale_pow2).to(output_ptr.dtype.element_ty),
|
||||
)
|
||||
|
||||
rope_offs = tl.arange(0, DIM_ROPE)
|
||||
bf16_off = (token_data_base + DIM_NOPE) // 2 + rope_offs
|
||||
rope_data = tl.load(buf_bf16_ptr + bf16_off)
|
||||
tl.store(
|
||||
output_ptr + out_row_base + DIM_NOPE + rope_offs,
|
||||
rope_data.to(output_ptr.dtype.element_ty),
|
||||
)
|
||||
|
||||
|
||||
def dequantize_k_cache_paged_ref(
|
||||
quant_k_cache: torch.Tensor,
|
||||
page_table_1_flattened: torch.Tensor,
|
||||
@@ -196,6 +382,29 @@ def dequantize_k_cache_paged_ref(
|
||||
return out
|
||||
|
||||
|
||||
def gather_dequant_requant_fp8_paged_ref(
|
||||
quant_k_cache: torch.Tensor,
|
||||
page_table_1_flattened: torch.Tensor,
|
||||
page_size: int,
|
||||
extra_rows: int = 0,
|
||||
) -> torch.Tensor:
|
||||
"""Torch reference for :func:`gather_dequant_requant_fp8_paged`."""
|
||||
active = dequantize_k_cache_paged_ref(
|
||||
quant_k_cache,
|
||||
page_table_1_flattened,
|
||||
page_size,
|
||||
).to(fp8_dtype)
|
||||
if extra_rows == 0:
|
||||
return active
|
||||
out = torch.zeros(
|
||||
(active.shape[0] + extra_rows, 1, DIM_NOPE + DIM_ROPE),
|
||||
dtype=fp8_dtype,
|
||||
device=active.device,
|
||||
)
|
||||
out[: active.shape[0]] = active
|
||||
return out
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
assert torch.cuda.is_available(), "this self-test needs a CUDA device"
|
||||
torch.manual_seed(0)
|
||||
|
||||
@@ -284,20 +284,116 @@ def sparse_mla_q8kv8_prefill_fwd(
|
||||
"""Run Q8KV8 (FP8) sparse prefill attention on SM90.
|
||||
|
||||
The kernel writes into three output tensors. By default fresh tensors
|
||||
are allocated and returned; callers that want to reuse buffers (e.g.
|
||||
for CUDA graph capture) may pass pre-allocated ``out`` / ``max_logits``
|
||||
/ ``lse`` tensors of the expected shape/dtype/device. The three output
|
||||
tensors must not alias each other.
|
||||
are allocated and returned; callers that want to reuse buffers may pass
|
||||
pre-allocated ``out`` / ``max_logits`` / ``lse`` tensors of the expected
|
||||
shape/dtype/device. The three output tensors must not alias each other.
|
||||
|
||||
Returns:
|
||||
out: [s_q, h_q, d_v], bfloat16
|
||||
max_logits: [s_q, h_q], float32
|
||||
lse: [s_q, h_q], float32
|
||||
"""
|
||||
# Validate ranks before unpacking shapes so malformed callers fail with a
|
||||
# clear error instead of a Python unpacking/indexing exception.
|
||||
if q.ndim != 3:
|
||||
raise ValueError(f"q must have shape (s_q, h_q, d_qk), got {tuple(q.shape)}")
|
||||
if kv.ndim != 3:
|
||||
raise ValueError(
|
||||
f"kv must have shape (s_kv, h_kv, d_qk), got {tuple(kv.shape)}"
|
||||
)
|
||||
if indices.ndim != 3:
|
||||
raise ValueError(
|
||||
"indices must have shape (s_q, h_kv, topk), " f"got {tuple(indices.shape)}"
|
||||
)
|
||||
|
||||
s_q, h_q, d_qk = q.shape
|
||||
s_kv = kv.shape[0]
|
||||
h_kv = kv.shape[1]
|
||||
s_kv, h_kv, kv_d_qk = kv.shape
|
||||
topk = indices.shape[2]
|
||||
device = q.device
|
||||
|
||||
# entry.cuh interprets q/kv as contiguous FP8 buffers and launches all
|
||||
# accesses on q's CUDA device. Reject contract violations before launch.
|
||||
if not q.is_cuda:
|
||||
raise ValueError("q must be a CUDA tensor")
|
||||
if not kv.is_cuda:
|
||||
raise ValueError("kv must be a CUDA tensor")
|
||||
if not indices.is_cuda:
|
||||
raise ValueError("indices must be a CUDA tensor")
|
||||
|
||||
if kv.device != device:
|
||||
raise ValueError(f"kv must be on q's device {device}, got {kv.device}")
|
||||
if indices.device != device:
|
||||
raise ValueError(
|
||||
f"indices must be on q's device {device}, got {indices.device}"
|
||||
)
|
||||
|
||||
if q.dtype != torch.float8_e4m3fn:
|
||||
raise ValueError(f"q must be torch.float8_e4m3fn, got {q.dtype}")
|
||||
if kv.dtype != torch.float8_e4m3fn:
|
||||
raise ValueError(f"kv must be torch.float8_e4m3fn, got {kv.dtype}")
|
||||
|
||||
if not q.is_contiguous():
|
||||
raise ValueError("q must be contiguous")
|
||||
if not kv.is_contiguous():
|
||||
raise ValueError("kv must be contiguous")
|
||||
if not indices.is_contiguous():
|
||||
raise ValueError("indices must be contiguous")
|
||||
|
||||
if kv_d_qk != d_qk:
|
||||
raise ValueError(f"kv d_qk must match q d_qk={d_qk}, got {kv_d_qk}")
|
||||
|
||||
# The CUDA implementation uses B_H=64 and launches h_q / B_H CTAs.
|
||||
# Reject unpadded TP-local head counts instead of launching zero CTAs and
|
||||
# returning uninitialized outputs, which can appear to callers as a hang or
|
||||
# a later collective failure.
|
||||
if h_q == 0 or h_q % 64 != 0:
|
||||
raise ValueError(
|
||||
"sparse_mla_q8kv8_prefill_fwd requires h_q padded to a positive "
|
||||
f"multiple of 64, got {h_q}"
|
||||
)
|
||||
|
||||
if h_kv != 1:
|
||||
raise ValueError(f"sparse_mla_q8kv8_prefill_fwd requires h_kv=1, got {h_kv}")
|
||||
|
||||
if d_qk not in (512, 576):
|
||||
raise ValueError(
|
||||
f"sparse_mla_q8kv8_prefill_fwd supports d_qk=512/576, got {d_qk}"
|
||||
)
|
||||
|
||||
if indices.shape[:2] != (s_q, h_kv):
|
||||
raise ValueError(
|
||||
"indices must have shape "
|
||||
f"({s_q}, {h_kv}, topk), got {tuple(indices.shape)}"
|
||||
)
|
||||
|
||||
if indices.dtype != torch.int32:
|
||||
raise ValueError(f"indices must be int32, got {indices.dtype}")
|
||||
|
||||
if topk == 0 or topk % 128 != 0:
|
||||
raise ValueError(
|
||||
"Q8KV8 sparse-prefill topk width must be a positive multiple of 128, "
|
||||
f"got {topk}"
|
||||
)
|
||||
|
||||
if topk_length is not None:
|
||||
if topk_length.shape != (s_q,) or topk_length.dtype != torch.int32:
|
||||
raise ValueError(
|
||||
f"topk_length must be int32 with shape ({s_q},), got "
|
||||
f"{tuple(topk_length.shape)}/{topk_length.dtype}"
|
||||
)
|
||||
if not topk_length.is_cuda:
|
||||
raise ValueError("topk_length must be a CUDA tensor")
|
||||
if topk_length.device != device:
|
||||
raise ValueError(
|
||||
"topk_length must be on q's device "
|
||||
f"{device}, got {topk_length.device}"
|
||||
)
|
||||
if not topk_length.is_contiguous():
|
||||
raise ValueError("topk_length must be contiguous")
|
||||
if torch.any(topk_length < 0).item() or torch.any(topk_length > topk).item():
|
||||
raise ValueError(
|
||||
"topk_length values must satisfy " f"0 <= topk_length <= topk ({topk})"
|
||||
)
|
||||
|
||||
if d_v != 512:
|
||||
raise ValueError(
|
||||
@@ -307,15 +403,49 @@ def sparse_mla_q8kv8_prefill_fwd(
|
||||
if attn_sink is not None and topk_length is None:
|
||||
raise ValueError("attn_sink requires topk_length to be provided as well")
|
||||
|
||||
device = q.device
|
||||
if attn_sink is not None:
|
||||
if attn_sink.shape != (h_q,) or attn_sink.dtype != torch.float32:
|
||||
raise ValueError(
|
||||
f"attn_sink must be float32 with shape ({h_q},), got "
|
||||
f"{tuple(attn_sink.shape)}/{attn_sink.dtype}"
|
||||
)
|
||||
if not attn_sink.is_cuda:
|
||||
raise ValueError("attn_sink must be a CUDA tensor")
|
||||
if attn_sink.device != device:
|
||||
raise ValueError(
|
||||
f"attn_sink must be on q's device {device}, got {attn_sink.device}"
|
||||
)
|
||||
if not attn_sink.is_contiguous():
|
||||
raise ValueError("attn_sink must be contiguous")
|
||||
|
||||
for name, scale in (("q_scale", q_scale), ("kv_scale", kv_scale)):
|
||||
if not isinstance(scale, torch.Tensor):
|
||||
raise ValueError(f"{name} must be a torch.Tensor")
|
||||
if not scale.is_cuda:
|
||||
raise ValueError(f"{name} must be a CUDA tensor")
|
||||
if scale.device != device:
|
||||
raise ValueError(
|
||||
f"{name} must be on q's device {device}, got {scale.device}"
|
||||
)
|
||||
if scale.dtype != torch.float32:
|
||||
raise ValueError(f"{name} must be float32, got {scale.dtype}")
|
||||
if scale.numel() != 1:
|
||||
raise ValueError(
|
||||
f"{name} must be a scalar tensor, got shape {tuple(scale.shape)}"
|
||||
)
|
||||
if not scale.is_contiguous():
|
||||
raise ValueError(f"{name} must be contiguous")
|
||||
|
||||
if out is None:
|
||||
out = torch.empty(s_q, h_q, d_v, dtype=torch.bfloat16, device=device)
|
||||
else:
|
||||
_check_out_buffer(out, "out", (s_q, h_q, d_v), torch.bfloat16, device)
|
||||
|
||||
if max_logits is None:
|
||||
max_logits = torch.empty(s_q, h_q, dtype=torch.float32, device=device)
|
||||
else:
|
||||
_check_out_buffer(max_logits, "max_logits", (s_q, h_q), torch.float32, device)
|
||||
|
||||
if lse is None:
|
||||
lse = torch.empty(s_q, h_q, dtype=torch.float32, device=device)
|
||||
else:
|
||||
|
||||
@@ -19,7 +19,11 @@ import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.kernels.ops.attention.dsv4.dequant_k_cache import (
|
||||
cast_q_fp8_for_q8kv8_prefill,
|
||||
dequantize_k_cache_paged,
|
||||
fp8_dtype,
|
||||
gather_dequant_requant_fp8_paged,
|
||||
q8kv8_padded_num_heads,
|
||||
)
|
||||
from sglang.kernels.ops.attention.dsv4.metadata_kernel import (
|
||||
init_compression_metadata as _init_compression_metadata_triton,
|
||||
@@ -56,8 +60,12 @@ from sglang.srt.layers.attention.dsv4.metadata import (
|
||||
from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
|
||||
SparsePrefillChunkCache,
|
||||
SparsePrefillWorkspace,
|
||||
use_dsv4_q8kv8_sparse_prefill,
|
||||
)
|
||||
from sglang.srt.layers.attention.verify_mask import (
|
||||
VerifyMask,
|
||||
maybe_create_verify_mask,
|
||||
)
|
||||
from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
|
||||
from sglang.srt.layers.cp.utils import is_cp_v2_active
|
||||
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||
@@ -70,7 +78,7 @@ from sglang.srt.speculative.ragged_verify import (
|
||||
read_ragged_verify_mode,
|
||||
resolve_ragged_verify_layout,
|
||||
)
|
||||
from sglang.srt.utils import ceil_align, is_cuda, is_xpu
|
||||
from sglang.srt.utils import ceil_align, is_cuda, is_sm90_supported, is_xpu
|
||||
from sglang.srt.utils.common import is_sm120_supported
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -552,6 +560,22 @@ class DeepseekV4AttnBackend(
|
||||
self.dsa_topk_backend: DSATopKBackend = DSATopKBackend(
|
||||
model_runner.server_args.dsa_topk_backend
|
||||
)
|
||||
self.dsv4_prefill_backend: str = getattr(
|
||||
model_runner.server_args, "dsv4_prefill_backend", "auto"
|
||||
)
|
||||
if use_dsv4_q8kv8_sparse_prefill(self.dsv4_prefill_backend):
|
||||
if not is_sm90_supported():
|
||||
raise ValueError(
|
||||
"DeepSeek-V4 flashmla_sparse_q8 prefill requires SM90 CUDA GPUs."
|
||||
)
|
||||
if self.head_dim_v != 512:
|
||||
raise ValueError(
|
||||
"DeepSeek-V4 flashmla_sparse_q8 prefill requires d_v=512, "
|
||||
f"got {self.head_dim_v}."
|
||||
)
|
||||
self._q8kv8_qpad_buf = None
|
||||
self._q8kv8_attn_sink_pad = None
|
||||
self._q8kv8_identity_scale = None
|
||||
self.topk = model_runner.server_args.speculative_eagle_topk or 0
|
||||
assert self.topk in [0, 1], "MTP Topk > 1 not supported for DeepSeek V4"
|
||||
self.mtp_enabled = self.topk > 0
|
||||
@@ -1673,6 +1697,16 @@ class DeepseekV4AttnBackend(
|
||||
or envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get()
|
||||
)
|
||||
):
|
||||
if use_dsv4_q8kv8_sparse_prefill(self.dsv4_prefill_backend):
|
||||
return self._forward_prefill_sparse_q8kv8(
|
||||
q=q,
|
||||
layer_id=layer_id,
|
||||
compress_ratio=compress_ratio,
|
||||
forward_batch=forward_batch,
|
||||
token_to_kv_pool=token_to_kv_pool,
|
||||
core_attn_metadata=core_attn_metadata,
|
||||
attn_sink=attn_sink,
|
||||
)
|
||||
return self._forward_prefill_sparse(
|
||||
q=q,
|
||||
layer_id=layer_id,
|
||||
@@ -1847,6 +1881,216 @@ class DeepseekV4AttnBackend(
|
||||
)
|
||||
return o
|
||||
|
||||
def _prepare_q8kv8_q_and_sink(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
attn_sink: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
|
||||
"""Pad TP-local heads to the SM90 kernel's 64-head CTA granularity."""
|
||||
num_tokens, num_heads, head_dim = q.shape
|
||||
padded_heads = q8kv8_padded_num_heads(num_heads)
|
||||
|
||||
qpad = getattr(self, "_q8kv8_qpad_buf", None)
|
||||
if (
|
||||
qpad is None
|
||||
or qpad.shape[0] < num_tokens
|
||||
or qpad.shape[1] != padded_heads
|
||||
or qpad.shape[2] != head_dim
|
||||
or qpad.device != q.device
|
||||
):
|
||||
qpad = torch.empty(
|
||||
(num_tokens, padded_heads, head_dim),
|
||||
dtype=fp8_dtype,
|
||||
device=q.device,
|
||||
)
|
||||
self._q8kv8_qpad_buf = qpad
|
||||
|
||||
qpad = qpad[:num_tokens]
|
||||
|
||||
q_fp8, _ = cast_q_fp8_for_q8kv8_prefill(
|
||||
q,
|
||||
padded_num_heads=padded_heads,
|
||||
out=qpad,
|
||||
)
|
||||
|
||||
sink_pad = getattr(self, "_q8kv8_attn_sink_pad", None)
|
||||
if (
|
||||
sink_pad is None
|
||||
or sink_pad.shape != (padded_heads,)
|
||||
or sink_pad.device != q.device
|
||||
):
|
||||
sink_pad = torch.zeros(padded_heads, dtype=torch.float32, device=q.device)
|
||||
self._q8kv8_attn_sink_pad = sink_pad
|
||||
|
||||
sink_pad[:num_heads].copy_(attn_sink.reshape(-1)[:num_heads])
|
||||
if padded_heads > num_heads:
|
||||
sink_pad[num_heads:].zero_()
|
||||
|
||||
scale = getattr(self, "_q8kv8_identity_scale", None)
|
||||
if scale is None or scale.device != q.device:
|
||||
scale = torch.ones((), dtype=torch.float32, device=q.device)
|
||||
self._q8kv8_identity_scale = scale
|
||||
|
||||
return q_fp8, sink_pad, scale, num_heads
|
||||
|
||||
def _forward_prefill_sparse_q8kv8(
|
||||
self,
|
||||
q: torch.Tensor,
|
||||
layer_id: int,
|
||||
compress_ratio: Literal[0, 4, 128],
|
||||
forward_batch: ForwardBatch,
|
||||
token_to_kv_pool: DeepSeekV4TokenToKVPool,
|
||||
core_attn_metadata: DSV4AttnMetadata,
|
||||
attn_sink: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Experimental DeepSeek-V4 sparse prefill path using Q8KV8 kernels.
|
||||
|
||||
This mirrors ``_forward_prefill_sparse``'s cache/index construction, but
|
||||
writes the gathered KV workspace as FP8 and calls the SM90 Q8KV8 sparse
|
||||
prefill kernel. The path is selected by ``--dsv4-prefill-backend
|
||||
flashmla_sparse_q8``; ``SGLANG_DSV4_Q8KV8_PREFILL`` remains as a debug
|
||||
override for focused runtime validation.
|
||||
"""
|
||||
|
||||
from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import (
|
||||
sparse_mla_q8kv8_prefill_fwd,
|
||||
)
|
||||
|
||||
q_flat = q.squeeze(1)
|
||||
if q_flat.ndim != 3:
|
||||
raise ValueError(
|
||||
f"Q8KV8 sparse prefill expects 3D Q after squeeze, got {q_flat.shape}"
|
||||
)
|
||||
|
||||
if attn_sink.numel() != q_flat.shape[1]:
|
||||
raise ValueError(
|
||||
f"attn_sink has {attn_sink.numel()} heads but Q has "
|
||||
f"{q_flat.shape[1]} local heads"
|
||||
)
|
||||
|
||||
q_fp8, attn_sink_pad, identity_scale, active_heads = (
|
||||
self._prepare_q8kv8_q_and_sink(q_flat, attn_sink)
|
||||
)
|
||||
|
||||
if not getattr(self, "_q8kv8_sparse_prefill_log_emitted", False):
|
||||
logger.info(
|
||||
"DSV4_Q8KV8_SPARSE_PREFILL_HIT layer_id=%s "
|
||||
"compress_ratio=%s q_shape=%s padded_heads=%s d_v=%s",
|
||||
layer_id,
|
||||
compress_ratio,
|
||||
tuple(q_flat.shape),
|
||||
q_fp8.shape[1],
|
||||
self.head_dim_v,
|
||||
)
|
||||
self._q8kv8_sparse_prefill_log_emitted = True
|
||||
|
||||
cache = self.forward_metadata.sparse_prefill_cache
|
||||
if cache is None:
|
||||
seq_lens_cpu = forward_batch.seq_lens_cpu
|
||||
assert seq_lens_cpu is not None
|
||||
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
|
||||
assert extend_seq_lens_cpu is not None
|
||||
total_swa = sum(
|
||||
min(int(seq_len), int(extend_len) + SWA_WINDOW - 1)
|
||||
for seq_len, extend_len in zip(
|
||||
seq_lens_cpu.tolist(), extend_seq_lens_cpu, strict=True
|
||||
)
|
||||
)
|
||||
cache = SparsePrefillChunkCache.build(
|
||||
seq_lens=forward_batch.seq_lens.to(torch.int32),
|
||||
extend_seq_lens=forward_batch.extend_seq_lens.to(torch.int32),
|
||||
req_pool_indices=forward_batch.req_pool_indices.to(torch.int32),
|
||||
req_to_token=self.req_to_token,
|
||||
full_to_swa=token_to_kv_pool.full_to_swa_index_mapping,
|
||||
swa_window_size=SWA_WINDOW,
|
||||
swa_page_size=token_to_kv_pool.swa_window_size,
|
||||
num_qo_tokens=q_flat.shape[0],
|
||||
max_seq_len=int(seq_lens_cpu.max().item()),
|
||||
total_swa=total_swa,
|
||||
)
|
||||
self.forward_metadata.sparse_prefill_cache = cache
|
||||
|
||||
compressed_slice = None
|
||||
extra_k_cache = None
|
||||
extra_page_size = None
|
||||
flat_token_ids = None
|
||||
|
||||
if compress_ratio == 0:
|
||||
workspace = self.sparse_prefill_workspace.get(
|
||||
cache.swa_token_ids.shape[0] + 1,
|
||||
dtype=fp8_dtype,
|
||||
)
|
||||
combined_indices = cache.c0_combined_indices
|
||||
combined_lens = cache.c0_combined_lens
|
||||
swa_slice = workspace
|
||||
else:
|
||||
extra_page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
|
||||
extra_k_cache = token_to_kv_pool.get_extra_key_buffer(layer_id)
|
||||
|
||||
if compress_ratio == 128:
|
||||
assert core_attn_metadata.c128_page_indices is not None
|
||||
cache.ensure_c128(core_attn_metadata.c128_page_indices)
|
||||
flat_token_ids = cache.c128_flat_token_ids
|
||||
combined_indices = cache.c128_combined_indices
|
||||
combined_lens = cache.c128_combined_lens
|
||||
else:
|
||||
assert core_attn_metadata.c4_sparse_raw_indices is not None, (
|
||||
"Q8KV8 sparse-prefill c4 path requires c4_sparse_raw_indices "
|
||||
"(allocated in init_flashmla_related when is_prefill=True)"
|
||||
)
|
||||
cache.ensure_c4(core_attn_metadata.page_table, extra_page_size)
|
||||
flat_token_ids = cache.c4_flat_token_ids
|
||||
combined_indices, combined_lens = cache.combine_c4_layer(
|
||||
c4_sparse_raw_indices=core_attn_metadata.c4_sparse_raw_indices[
|
||||
: cache.num_qo_tokens
|
||||
],
|
||||
)
|
||||
|
||||
n_compressed = flat_token_ids.shape[0]
|
||||
workspace = self.sparse_prefill_workspace.get(
|
||||
n_compressed + cache.swa_token_ids.shape[0] + 1,
|
||||
dtype=fp8_dtype,
|
||||
)
|
||||
compressed_slice = workspace[:n_compressed]
|
||||
swa_slice = workspace[n_compressed:]
|
||||
|
||||
if compressed_slice is not None:
|
||||
gather_dequant_requant_fp8_paged(
|
||||
extra_k_cache,
|
||||
flat_token_ids,
|
||||
page_size=extra_page_size,
|
||||
out=compressed_slice,
|
||||
)
|
||||
|
||||
gather_dequant_requant_fp8_paged(
|
||||
token_to_kv_pool.get_swa_key_buffer_radix(layer_id),
|
||||
cache.swa_token_ids,
|
||||
page_size=cache.swa_page_size,
|
||||
extra_rows=1,
|
||||
out=swa_slice,
|
||||
)
|
||||
|
||||
sentinel_row = workspace.shape[0] - 1
|
||||
q8_indices = torch.where(
|
||||
combined_indices < 0,
|
||||
torch.full_like(combined_indices, sentinel_row),
|
||||
combined_indices,
|
||||
)
|
||||
|
||||
o, _, _ = sparse_mla_q8kv8_prefill_fwd(
|
||||
q=q_fp8,
|
||||
kv=workspace,
|
||||
indices=q8_indices.unsqueeze(1),
|
||||
sm_scale=self.softmax_scale,
|
||||
q_scale=identity_scale,
|
||||
kv_scale=identity_scale,
|
||||
d_v=self.head_dim_v,
|
||||
attn_sink=attn_sink_pad,
|
||||
topk_length=combined_lens,
|
||||
)
|
||||
|
||||
return o[:, :active_heads]
|
||||
|
||||
def expand_prefill_casually(
|
||||
self,
|
||||
num_tokens: int,
|
||||
|
||||
@@ -32,6 +32,7 @@ For SWA-only layers callers pass ``topk=0``, ``compressed_base = 0`` (the
|
||||
compressed branch becomes a no-op) and any ``compress_ratio >= 1``.
|
||||
"""
|
||||
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
@@ -47,7 +48,8 @@ SPARSE_PREFILL_TOPK_ALIGNMENT = 128
|
||||
# Bf16 workspace per-token width, matching ``dequantize_k_cache_paged``'s
|
||||
# output: 448 fp8 nope (dequanted) + 64 bf16 rope = 512.
|
||||
WORKSPACE_DIM = DIM_NOPE + DIM_ROPE
|
||||
|
||||
DSV4_Q8KV8_PREFILL_ENV = "SGLANG_DSV4_Q8KV8_PREFILL"
|
||||
DSV4_Q8KV8_PREFILL_LOG_ENV = "SGLANG_DSV4_Q8KV8_PREFILL_LOG"
|
||||
|
||||
from sglang.kernels.ops.attention.dsv4.sparse_prefill_kernels import (
|
||||
_build_swa_token_ids_kernel,
|
||||
@@ -55,6 +57,24 @@ from sglang.kernels.ops.attention.dsv4.sparse_prefill_kernels import (
|
||||
)
|
||||
|
||||
|
||||
def use_dsv4_q8kv8_sparse_prefill(dsv4_prefill_backend: str = "auto") -> bool:
|
||||
"""Return whether DeepSeek-V4 sparse prefill should use Q8KV8.
|
||||
|
||||
``dsv4_prefill_backend`` is the production configuration. The environment
|
||||
variable remains as a debug override while the runtime path is being
|
||||
hardened: truthy values force Q8 on, falsy values force it off.
|
||||
"""
|
||||
env_value = os.getenv(DSV4_Q8KV8_PREFILL_ENV)
|
||||
if env_value is not None:
|
||||
return env_value.lower() in {
|
||||
"1",
|
||||
"true",
|
||||
"yes",
|
||||
"on",
|
||||
}
|
||||
return dsv4_prefill_backend == "flashmla_sparse_q8"
|
||||
|
||||
|
||||
class SparsePrefillWorkspace:
|
||||
"""Backend-owned scratch storage for sparse prefill KV dequantization.
|
||||
|
||||
@@ -68,13 +88,18 @@ class SparsePrefillWorkspace:
|
||||
self.device = device
|
||||
self._buffer: Optional[torch.Tensor] = None
|
||||
|
||||
def get(self, num_tokens: int) -> torch.Tensor:
|
||||
def get(
|
||||
self,
|
||||
num_tokens: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> torch.Tensor:
|
||||
assert num_tokens > 0
|
||||
current_capacity = self._buffer.shape[0] if self._buffer is not None else 0
|
||||
if num_tokens > current_capacity:
|
||||
current_dtype = self._buffer.dtype if self._buffer is not None else None
|
||||
if num_tokens > current_capacity or dtype != current_dtype:
|
||||
self._buffer = torch.empty(
|
||||
(num_tokens, 1, WORKSPACE_DIM),
|
||||
dtype=torch.bfloat16,
|
||||
dtype=dtype,
|
||||
device=self.device,
|
||||
)
|
||||
return self._buffer[:num_tokens]
|
||||
|
||||
@@ -359,6 +359,12 @@ DSA_CHOICES = [
|
||||
]
|
||||
NSA_CHOICES = DSA_CHOICES # deprecated alias
|
||||
|
||||
DSV4_PREFILL_BACKEND_CHOICES = [
|
||||
"auto",
|
||||
"flashmla_sparse",
|
||||
"flashmla_sparse_q8",
|
||||
]
|
||||
|
||||
DSA_TOPK_BACKEND_CHOICES = ["sgl-kernel", "torch", "flashinfer"]
|
||||
|
||||
DSA_PAGED_MQA_LOGITS_BACKEND_CHOICES = ["auto", "deepgemm", "cutedsl", "aiter"]
|
||||
@@ -1805,6 +1811,18 @@ class ServerArgs:
|
||||
),
|
||||
NS("exec.kernel"),
|
||||
] = None
|
||||
dsv4_prefill_backend: A[
|
||||
str,
|
||||
Arg(
|
||||
help=(
|
||||
"DeepSeek-V4 sparse prefill backend. 'auto' and "
|
||||
"'flashmla_sparse' use the existing BF16 sparse prefill path; "
|
||||
"'flashmla_sparse_q8' enables the Q8KV8 sparse prefill path."
|
||||
),
|
||||
choices=DSV4_PREFILL_BACKEND_CHOICES,
|
||||
),
|
||||
NS("exec.kernel"),
|
||||
] = "auto"
|
||||
dsa_decode_backend: A[
|
||||
Optional[str],
|
||||
Arg(
|
||||
|
||||
Reference in New Issue
Block a user