From 9db4ba8da166ce62e61b9d4f471f53065cdf2874 Mon Sep 17 00:00:00 2001 From: shiyang814-cpu Date: Thu, 20 Aug 2026 10:23:49 +0800 Subject: [PATCH] [DeepSeek-V4] Add Q8KV8 sparse MLA prefill runtime backend (#32327) Co-authored-by: Ho-Ren (Jack) Chuang Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- .../ops/attention/dsv4/dequant_k_cache.py | 209 ++++++ .../sparse_mla_q8kv8_prefill_sm90.py | 144 +++- .../layers/attention/deepseek_v4_backend.py | 248 ++++++- .../attention/dsv4/sparse_prefill_utils.py | 33 +- python/sglang/srt/server_args.py | 18 + .../test_q8kv8_sparse_prefill_backend.py | 681 ++++++++++++++++++ .../unit/server_args/test_server_args.py | 17 + 7 files changed, 1337 insertions(+), 13 deletions(-) create mode 100644 test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py diff --git a/python/sglang/kernels/ops/attention/dsv4/dequant_k_cache.py b/python/sglang/kernels/ops/attention/dsv4/dequant_k_cache.py index 01ad58aad..ea62cc7a8 100644 --- a/python/sglang/kernels/ops/attention/dsv4/dequant_k_cache.py +++ b/python/sglang/kernels/ops/attention/dsv4/dequant_k_cache.py @@ -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) diff --git a/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py b/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py index 734a78c63..bd014c0f8 100644 --- a/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py +++ b/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py @@ -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: diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 37081e06e..8b1853441 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -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, diff --git a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py index 88ba1ff72..3ca85e5f2 100644 --- a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py +++ b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py @@ -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] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e0e0c53bc..679f56bf6 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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( diff --git a/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py b/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py new file mode 100644 index 000000000..ca2adb70e --- /dev/null +++ b/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py @@ -0,0 +1,681 @@ +"""DeepSeek-V4 Q8KV8 sparse-prefill backend helper tests. + +These tests avoid starting a full server. They construct the minimum V4 +metadata and token-pool surface consumed by the sparse-prefill helpers, then +compare the BF16 sparse path's gathered workspace against the Q8 path's FP8 +workspace after dequantizing it back to BF16. +""" + +from __future__ import annotations + +import sys +import types +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +from sglang.kernels.ops.attention.dsv4.index_buf_accessor import SetKAndS +from sglang.kernels.ops.attention.dsv4.quant_k_cache import ( + quant_to_nope_fp8_rope_bf16_pack_triton, +) +from sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90 import ( + sparse_mla_q8kv8_prefill_fwd, +) +from sglang.srt.layers.attention.deepseek_v4_backend import DeepseekV4AttnBackend +from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import ( + SparsePrefillChunkCache, + SparsePrefillWorkspace, + use_dsv4_q8kv8_sparse_prefill, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.utils import is_sm90_supported +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large") + + +def test_q8kv8_sparse_prefill_backend_selector_uses_cli_value(): + assert not use_dsv4_q8kv8_sparse_prefill() + assert not use_dsv4_q8kv8_sparse_prefill("auto") + assert not use_dsv4_q8kv8_sparse_prefill("flashmla_sparse") + assert use_dsv4_q8kv8_sparse_prefill("flashmla_sparse_q8") + + +class _Pool: + def __init__(self, page_size: int): + self.page_size = page_size + + +class _Capture: + def __init__(self): + self.calls = [] + + def record(self, **kwargs): + cloned = {} + for name, value in kwargs.items(): + if isinstance(value, torch.Tensor): + cloned[name] = value.detach().clone() + else: + cloned[name] = value + self.calls.append(cloned) + + +class _TokenToKVPool: + def __init__( + self, + *, + swa_key_buffer: torch.Tensor, + full_to_swa_index_mapping: torch.Tensor, + page_size: int, + extra_key_buffer: torch.Tensor | None = None, + ): + self._swa_key_buffer = swa_key_buffer + self._extra_key_buffer = ( + extra_key_buffer if extra_key_buffer is not None else swa_key_buffer + ) + self.full_to_swa_index_mapping = full_to_swa_index_mapping + self.swa_window_size = page_size + + def get_swa_key_buffer_radix(self, layer_id: int) -> torch.Tensor: + _ = layer_id + return self._swa_key_buffer + + def get_extra_key_page_size(self, layer_id: int) -> int: + _ = layer_id + return self.swa_window_size + + def get_extra_key_buffer(self, layer_id: int) -> torch.Tensor: + _ = layer_id + return self._extra_key_buffer + + +def _sm90_available() -> bool: + return torch.cuda.is_available() and is_sm90_supported() + + +def _make_v4_paged_kv_cache( + *, + total_slots: int, + page_size: int, + seed: int, + device: torch.device, +) -> torch.Tensor: + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + num_pages = (total_slots + page_size - 1) // page_size + total_slots = num_pages * page_size + bytes_per_token = 448 + 64 * 2 + 8 + quant_k_cache = torch.zeros( + num_pages, + page_size * bytes_per_token, + dtype=torch.uint8, + device=device, + ) + + k_bf16 = (torch.randn(total_slots, 512, device=device) * 0.25).to(torch.bfloat16) + pack = quant_to_nope_fp8_rope_bf16_pack_triton(k_bf16) + loc = torch.arange(total_slots, dtype=torch.int32, device=device) + SetKAndS.torch(_Pool(page_size), quant_k_cache, loc, pack) + return quant_k_cache + + +def _make_forward_batch_and_mapping( + device: torch.device, +) -> tuple[ForwardBatch, torch.Tensor]: + seq_lens = torch.tensor([96, 144], dtype=torch.int32, device=device) + extend_seq_lens = torch.tensor([3, 2], dtype=torch.int32, device=device) + req_pool_indices = torch.tensor([0, 1], dtype=torch.int32, device=device) + seq0 = int(seq_lens[0].item()) + seq1 = int(seq_lens[1].item()) + + req_to_token = torch.zeros( + (2, int(seq_lens.max().item())), dtype=torch.int32, device=device + ) + req_to_token[0, :seq0] = torch.arange(seq0, dtype=torch.int32, device=device) + req1_base = 192 + req_to_token[1, :seq1] = req1_base + torch.arange( + seq1, dtype=torch.int32, device=device + ) + + forward_batch = ForwardBatch( + forward_mode=ForwardMode.EXTEND, + batch_size=2, + input_ids=torch.zeros( + int(extend_seq_lens.sum().item()), dtype=torch.int32, device=device + ), + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + out_cache_loc=torch.zeros( + int(extend_seq_lens.sum().item()), dtype=torch.int32, device=device + ), + seq_lens_sum=int(seq_lens.sum().item()), + seq_lens_cpu=seq_lens.detach().cpu(), + extend_num_tokens=int(extend_seq_lens.sum().item()), + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=[int(x) for x in extend_seq_lens.detach().cpu().tolist()], + ) + return forward_batch, req_to_token + + +def _make_backend( + device: torch.device, + req_to_token: torch.Tensor, + dsv4_prefill_backend: str = "auto", +) -> DeepseekV4AttnBackend: + backend = DeepseekV4AttnBackend.__new__(DeepseekV4AttnBackend) + backend.forward_metadata = SimpleNamespace(sparse_prefill_cache=None) + backend.req_to_token = req_to_token + backend.sparse_prefill_workspace = SparsePrefillWorkspace(device) + backend.softmax_scale = 512**-0.5 + backend.head_dim_v = 512 + backend.dsv4_prefill_backend = dsv4_prefill_backend + return backend + + +def _make_sparse_prefill_case( + device: torch.device, + local_heads: int = 64, +): + page_size = 64 + total_slots = 384 + forward_batch, req_to_token = _make_forward_batch_and_mapping(device) + backend = _make_backend(device, req_to_token) + quant_k_cache = _make_v4_paged_kv_cache( + total_slots=total_slots, + page_size=page_size, + seed=3, + device=device, + ) + extra_k_cache = _make_v4_paged_kv_cache( + total_slots=total_slots, + page_size=page_size, + seed=7, + device=device, + ) + token_to_kv_pool = _TokenToKVPool( + swa_key_buffer=quant_k_cache, + extra_key_buffer=extra_k_cache, + full_to_swa_index_mapping=torch.arange( + total_slots, dtype=torch.int64, device=device + ), + page_size=page_size, + ) + + generator = torch.Generator(device=device) + generator.manual_seed(11) + q = ( + torch.randn( + forward_batch.extend_num_tokens, + 1, + local_heads, + 512, + device=device, + generator=generator, + ) + * 0.05 + ).to(torch.bfloat16) + attn_sink = torch.zeros(local_heads, dtype=torch.float32, device=device) + core_attn_metadata = SimpleNamespace() + return backend, forward_batch, token_to_kv_pool, q, attn_sink, core_attn_metadata + + +def _populate_compress_metadata( + core_attn_metadata: SimpleNamespace, + *, + compress_ratio: int, + device: torch.device, +) -> None: + if compress_ratio == 4: + core_attn_metadata.page_table = torch.zeros( + (2, 4), dtype=torch.int32, device=device + ) + core_attn_metadata.c4_sparse_raw_indices = torch.zeros( + (16, 1), dtype=torch.int32, device=device + ) + elif compress_ratio == 128: + core_attn_metadata.c128_page_indices = torch.zeros( + (16, 1), dtype=torch.int32, device=device + ) + + +@contextmanager +def _patched_compressed_sparse_cache_paths(compress_ratio: int): + if compress_ratio == 0: + yield + return + + old_ensure_c4 = SparsePrefillChunkCache.ensure_c4 + old_ensure_c128 = SparsePrefillChunkCache.ensure_c128 + old_combine_c4_layer = SparsePrefillChunkCache.combine_c4_layer + + def _with_compressed_prefix(cache: SparsePrefillChunkCache, n_compressed: int): + shifted_swa = torch.where( + cache.c0_combined_indices >= 0, + cache.c0_combined_indices + n_compressed, + cache.c0_combined_indices, + ) + n_prefix = min(n_compressed, shifted_swa.shape[1]) + if n_prefix > 0: + shifted_swa[:, :n_prefix] = torch.arange( + n_prefix, dtype=shifted_swa.dtype, device=shifted_swa.device + ) + combined_lens = torch.clamp( + cache.c0_combined_lens + n_prefix, + max=shifted_swa.shape[1], + ) + return shifted_swa, combined_lens + + def fake_ensure_c128(self, c128_page_indices): + _ = c128_page_indices + n_compressed = 8 + self.c128_flat_token_ids = torch.arange( + n_compressed, dtype=torch.int64, device=self.swa_token_ids.device + ) + self.c128_combined_indices, self.c128_combined_lens = _with_compressed_prefix( + self, n_compressed + ) + + def fake_ensure_c4(self, page_table, extra_page_size): + _ = page_table, extra_page_size + n_compressed = 8 + self.c4_flat_token_ids = torch.arange( + n_compressed, dtype=torch.int64, device=self.swa_token_ids.device + ) + + def fake_combine_c4_layer(self, c4_sparse_raw_indices): + _ = c4_sparse_raw_indices + return _with_compressed_prefix(self, self.c4_flat_token_ids.shape[0]) + + SparsePrefillChunkCache.ensure_c128 = fake_ensure_c128 + SparsePrefillChunkCache.ensure_c4 = fake_ensure_c4 + SparsePrefillChunkCache.combine_c4_layer = fake_combine_c4_layer + try: + yield + finally: + SparsePrefillChunkCache.ensure_c4 = old_ensure_c4 + SparsePrefillChunkCache.ensure_c128 = old_ensure_c128 + SparsePrefillChunkCache.combine_c4_layer = old_combine_c4_layer + + +def _make_q8kv8_kernel_args( + *, + device: torch.device, + s_q: int = 4, + h_q: int = 64, + d_qk: int = 512, + s_kv: int = 256, + h_kv: int = 1, + topk: int = 128, +): + q = (torch.randn(s_q, h_q, d_qk, device=device) * 0.05).to(torch.float8_e4m3fn) + kv = (torch.randn(s_kv, h_kv, d_qk, device=device) * 0.05).to(torch.float8_e4m3fn) + indices = torch.randint( + 0, s_kv, (s_q, h_kv, topk), dtype=torch.int32, device=device + ) + topk_length = torch.full((s_q,), topk, dtype=torch.int32, device=device) + return { + "q": q.contiguous(), + "kv": kv.contiguous(), + "indices": indices.contiguous(), + "sm_scale": 512**-0.5, + "q_scale": torch.ones((), dtype=torch.float32, device=device), + "kv_scale": torch.ones((), dtype=torch.float32, device=device), + "d_v": 512, + "attn_sink": torch.zeros(h_q, dtype=torch.float32, device=device), + "topk_length": topk_length, + } + + +@contextmanager +def _patched_sparse_kernels( + bf16_capture: _Capture, + q8_capture: _Capture, +): + def fake_flash_mla_sparse_fwd( + *, + q, + kv, + indices, + sm_scale, + d_v, + attn_sink, + topk_length, + ): + bf16_capture.record( + q=q, + kv=kv, + indices=indices, + sm_scale=sm_scale, + d_v=d_v, + attn_sink=attn_sink, + topk_length=topk_length, + ) + out = torch.zeros( + (q.shape[0], q.shape[1], d_v), dtype=torch.bfloat16, device=q.device + ) + meta = torch.zeros( + (q.shape[0], q.shape[1]), dtype=torch.float32, device=q.device + ) + return out, meta, meta + + def fake_sparse_mla_q8kv8_prefill_fwd( + *, + q, + kv, + indices, + sm_scale, + q_scale, + kv_scale, + d_v, + attn_sink, + topk_length, + ): + q8_capture.record( + q=q, + kv=kv, + indices=indices, + sm_scale=sm_scale, + q_scale=q_scale, + kv_scale=kv_scale, + d_v=d_v, + attn_sink=attn_sink, + topk_length=topk_length, + ) + out = torch.zeros( + (q.shape[0], q.shape[1], d_v), dtype=torch.bfloat16, device=q.device + ) + meta = torch.zeros( + (q.shape[0], q.shape[1]), dtype=torch.float32, device=q.device + ) + return out, meta, meta + + sgl_kernel_pkg = types.ModuleType("sgl_kernel") + flash_mla_mod = types.ModuleType("sgl_kernel.flash_mla") + flash_mla_mod.flash_mla_sparse_fwd = fake_flash_mla_sparse_fwd + sgl_kernel_pkg.flash_mla = flash_mla_mod + + q8_module_name = "sglang.kernels.ops.attention.sparse_mla_q8kv8_prefill_sm90" + q8_mod = types.ModuleType(q8_module_name) + q8_mod.sparse_mla_q8kv8_prefill_fwd = fake_sparse_mla_q8kv8_prefill_fwd + + old_modules = { + name: sys.modules.get(name) + for name in ( + "sgl_kernel", + "sgl_kernel.flash_mla", + q8_module_name, + ) + } + sys.modules["sgl_kernel"] = sgl_kernel_pkg + sys.modules["sgl_kernel.flash_mla"] = flash_mla_mod + sys.modules[q8_module_name] = q8_mod + try: + yield + finally: + for name, old_value in old_modules.items(): + if old_value is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = old_value + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +@pytest.mark.parametrize("compress_ratio", [0, 4, 128]) +def test_q8kv8_sparse_prefill_helper_builds_fp8_workspace_matching_bf16_path( + compress_ratio: int, +): + from sglang.kernels.ops.attention.dsv4.dequant_k_cache import fp8_dtype + + device = torch.device("cuda") + backend, forward_batch, token_to_kv_pool, q, attn_sink, core_attn_metadata = ( + _make_sparse_prefill_case(device, local_heads=16) + ) + _populate_compress_metadata( + core_attn_metadata, + compress_ratio=compress_ratio, + device=device, + ) + + bf16_capture = _Capture() + q8_capture = _Capture() + with _patched_sparse_kernels( + bf16_capture, q8_capture + ), _patched_compressed_sparse_cache_paths(compress_ratio): + bf16_out = backend._forward_prefill_sparse( + q=q, + layer_id=0, + 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, + ) + sparse_cache = backend.forward_metadata.sparse_prefill_cache + q8_out = backend._forward_prefill_sparse_q8kv8( + q=q, + layer_id=0, + 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, + ) + + assert backend.forward_metadata.sparse_prefill_cache is sparse_cache + assert ( + bf16_out.shape + == q8_out.shape + == ( + forward_batch.extend_num_tokens, + 16, + 512, + ) + ) + assert len(bf16_capture.calls) == 1 + assert len(q8_capture.calls) == 1 + + bf16_call = bf16_capture.calls[0] + q8_call = q8_capture.calls[0] + bf16_kv = bf16_call["kv"] + q8_kv = q8_call["kv"] + + assert bf16_kv.dtype == torch.bfloat16 + assert q8_kv.dtype == fp8_dtype + assert q8_call["q"].dtype == fp8_dtype + assert q8_call["q"].shape[1] == 64 + assert torch.count_nonzero(q8_call["q"][:, 16:]).item() == 0 + assert q8_call["attn_sink"].shape == (64,) + assert q8_kv.shape[0] == bf16_kv.shape[0] + 1 + torch.testing.assert_close( + q8_kv[:-1].to(torch.bfloat16).float(), + bf16_kv.float(), + atol=3e-2, + rtol=2e-1, + ) + assert torch.equal( + q8_kv[-1].to(torch.bfloat16), + torch.zeros_like(q8_kv[-1].to(torch.bfloat16)), + ) + + bf16_indices = bf16_call["indices"] + q8_indices = q8_call["indices"] + sentinel_row = q8_kv.shape[0] - 1 + valid_mask = bf16_indices >= 0 + assert torch.equal(q8_indices[valid_mask], bf16_indices[valid_mask]) + assert torch.equal( + q8_indices[~valid_mask], + torch.full_like(q8_indices[~valid_mask], sentinel_row), + ) + assert torch.equal(q8_call["topk_length"], bf16_call["topk_length"]) + assert q8_call["q_scale"].item() == pytest.approx(1.0) + assert q8_call["kv_scale"].item() == pytest.approx(1.0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +def test_q8kv8_sparse_prefill_rejects_topk_64_before_cuda_launch(): + args = _make_q8kv8_kernel_args(device=torch.device("cuda"), topk=64) + + with pytest.raises(ValueError, match="positive multiple of 128"): + sparse_mla_q8kv8_prefill_fwd(**args) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +@pytest.mark.parametrize( + ("mutate", "error_match"), + [ + ( + lambda args: args.update(q=args["q"].float()), + "q must be torch.float8_e4m3fn", + ), + ( + lambda args: args.update(kv=args["kv"].float()), + "kv must be torch.float8_e4m3fn", + ), + ( + lambda args: args.update( + q=torch.empty( + args["q"].shape[0], + args["q"].shape[1], + args["q"].shape[2] + 1, + dtype=args["q"].dtype, + device=args["q"].device, + )[:, :, : args["q"].shape[2]] + ), + "q must be contiguous", + ), + ( + lambda args: args.update( + q_scale=torch.ones(2, dtype=torch.float32, device=args["q"].device) + ), + "q_scale must be a scalar tensor", + ), + ( + lambda args: args.update( + kv_scale=torch.ones((), dtype=torch.float16, device=args["q"].device) + ), + "kv_scale must be float32", + ), + ], +) +def test_q8kv8_sparse_prefill_rejects_invalid_tensor_contracts( + mutate, + error_match: str, +): + args = _make_q8kv8_kernel_args(device=torch.device("cuda"), topk=128) + mutate(args) + + with pytest.raises(ValueError, match=error_match): + sparse_mla_q8kv8_prefill_fwd(**args) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available") +@pytest.mark.parametrize("bad_length", [-1, 129]) +def test_q8kv8_sparse_prefill_rejects_invalid_topk_length_bounds( + bad_length: int, +): + args = _make_q8kv8_kernel_args(device=torch.device("cuda"), topk=128) + args["topk_length"][0] = bad_length + + with pytest.raises(ValueError, match="0 <= topk_length <= topk"): + sparse_mla_q8kv8_prefill_fwd(**args) + + +@pytest.mark.skipif( + not _sm90_available(), reason="Q8KV8 sparse prefill requires SM90 CUDA" +) +def test_q8kv8_sparse_prefill_real_kernel_matches_bf16_sparse_path(): + device = torch.device("cuda") + backend, forward_batch, token_to_kv_pool, q, attn_sink, core_attn_metadata = ( + _make_sparse_prefill_case(device, local_heads=64) + ) + + bf16_out = backend._forward_prefill_sparse( + q=q, + layer_id=0, + compress_ratio=0, + forward_batch=forward_batch, + token_to_kv_pool=token_to_kv_pool, + core_attn_metadata=core_attn_metadata, + attn_sink=attn_sink, + ) + sparse_cache = backend.forward_metadata.sparse_prefill_cache + q8_out = backend._forward_prefill_sparse_q8kv8( + q=q, + layer_id=0, + compress_ratio=0, + forward_batch=forward_batch, + token_to_kv_pool=token_to_kv_pool, + core_attn_metadata=core_attn_metadata, + attn_sink=attn_sink, + ) + torch.cuda.synchronize() + + assert backend.forward_metadata.sparse_prefill_cache is sparse_cache + assert ( + bf16_out.shape + == q8_out.shape + == ( + forward_batch.extend_num_tokens, + 64, + 512, + ) + ) + assert bf16_out.dtype == torch.bfloat16 + assert q8_out.dtype == torch.bfloat16 + assert torch.isfinite(bf16_out.float()).all() + assert torch.isfinite(q8_out.float()).all() + + abs_diff = (q8_out.float() - bf16_out.float()).abs() + assert abs_diff.mean().item() < 0.03 + assert torch.quantile(abs_diff.flatten(), 0.99).item() < 0.2 + torch.testing.assert_close( + q8_out.float(), + bf16_out.float(), + atol=2.5e-1, + rtol=3.0e-1, + ) + + +@pytest.mark.skipif( + not _sm90_available(), reason="Q8KV8 sparse prefill requires SM90 CUDA" +) +def test_q8kv8_sparse_prefill_real_kernel_repeated_launch_stable(): + args = _make_q8kv8_kernel_args( + device=torch.device("cuda"), + s_q=512, + h_q=64, + d_qk=512, + s_kv=1024, + h_kv=1, + topk=256, + ) + + baseline = None + for _ in range(10): + out, max_logits, lse = sparse_mla_q8kv8_prefill_fwd(**args) + torch.cuda.synchronize() + + assert out.shape == (512, 64, 512) + assert max_logits.shape == (512, 64) + assert lse.shape == (512, 64) + assert torch.isfinite(out.float()).all() + assert torch.isfinite(max_logits).all() + assert torch.isfinite(lse).all() + + current = out.float().detach().clone() + if baseline is None: + baseline = current + else: + torch.testing.assert_close( + current, + baseline, + atol=1e-2, + rtol=1e-2, + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 78b1a8368..28ea5f7ed 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -52,6 +52,23 @@ class TestPrepareServerArgs(CustomTestCase): ): ServerArgs(model_path="dummy", prefill_decode_interval=-1) + def test_dsv4_prefill_backend_cli_choices(self): + parser = server_args_module.argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + + base_args = ["--model-path", "dummy-model"] + + default_args = parser.parse_args(base_args) + self.assertEqual(default_args.dsv4_prefill_backend, "auto") + + q8_args = parser.parse_args( + base_args + ["--dsv4-prefill-backend", "flashmla_sparse_q8"] + ) + self.assertEqual(q8_args.dsv4_prefill_backend, "flashmla_sparse_q8") + + with self.assertRaises(SystemExit): + parser.parse_args(base_args + ["--dsv4-prefill-backend", "flashmla_kv"]) + def test_return_hidden_states_mode_configuration(self): disabled = ServerArgs(model_path="dummy") self.assertFalse(disabled.enable_return_hidden_states)