315 lines
11 KiB
Python
315 lines
11 KiB
Python
"""SM120 fp8_paged_mqa_logits_torch_sm120 vectorized PyTorch fallback tests.
|
|
|
|
Validates that the vectorized SM120-specific implementation matches the loopy
|
|
reference (`fp8_paged_mqa_logits_torch`) and is CUDA-graph compatible.
|
|
|
|
Coverage:
|
|
- Numeric equivalence vs reference at small shapes
|
|
- Both KV-cache dtype views (uint8 raw / float8_e4m3fn) — guards against the
|
|
historic garbled-output bug where Triton kernels treated uint8 bytes as raw
|
|
integers instead of FP8 (dsv4_sm120_progress.md §4)
|
|
- Variable per-batch seq_lens with -inf masking semantics
|
|
- CUDA graph capture + replay equivalence (no .item() / data-dependent shapes)
|
|
- Shape-assertion guards
|
|
|
|
Pure-PyTorch implementations on any CUDA GPU — no SM120 hardware required.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import unittest
|
|
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.dsv4.indexer import (
|
|
FP8_DTYPE,
|
|
fp8_paged_mqa_logits_torch,
|
|
fp8_paged_mqa_logits_torch_sm120,
|
|
)
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small")
|
|
|
|
|
|
# DSv4 indexer cache layout (fixed by deepseek_v4_memory_pool.DeepSeekV4IndexerPool):
|
|
# page_size = 64 tokens
|
|
# head_dim = 128 (FP8 values per token)
|
|
# quant_block_size = 128 -> num_scales_per_token = 1 (fp32 scale)
|
|
# per-page memory: [page_size*head_dim FP8 bytes][page_size*4 scale bytes]
|
|
# = [8192][256] = 8448 bytes
|
|
# The (block_size, 1, head_dim+4) shape is a fiction for downstream consumers.
|
|
PAGE_SIZE = 64
|
|
HEAD_DIM = 128
|
|
SCALE_BYTES_PER_TOKEN = 4
|
|
HEAD_DIM_WITH_SF = HEAD_DIM + SCALE_BYTES_PER_TOKEN # 132
|
|
PAGE_BYTES = PAGE_SIZE * HEAD_DIM + PAGE_SIZE * SCALE_BYTES_PER_TOKEN # 8448
|
|
|
|
|
|
def _build_kvcache(
|
|
num_pages: int,
|
|
*,
|
|
dtype_view: torch.dtype,
|
|
device: torch.device,
|
|
seed: int = 0,
|
|
) -> torch.Tensor:
|
|
"""Construct a paged KV cache matching the production layout.
|
|
|
|
Returns a tensor shaped (num_pages, PAGE_SIZE, 1, HEAD_DIM_WITH_SF) whose
|
|
underlying memory is the same regardless of `dtype_view`:
|
|
- [0 : PAGE_SIZE*HEAD_DIM) : random FP8 bit patterns (values)
|
|
- [PAGE_SIZE*HEAD_DIM : PAGE_BYTES) : random positive fp32 scales (bytes)
|
|
"""
|
|
g = torch.Generator(device="cpu").manual_seed(seed)
|
|
raw = torch.empty(num_pages, PAGE_BYTES, dtype=torch.uint8, device=device)
|
|
|
|
# Random FP8 byte pattern for the value section. Bias away from extreme
|
|
# bit patterns that map to NaN/Inf in float8_e4m3fn (sign|exp4|mantissa3;
|
|
# exp=0xF mantissa!=0 -> NaN). Restricting to [0, 0x6F] keeps |x| < 256.
|
|
val_bytes = torch.randint(
|
|
0, 0x70, (num_pages, PAGE_SIZE * HEAD_DIM), generator=g, dtype=torch.uint8
|
|
).to(device)
|
|
raw[:, : PAGE_SIZE * HEAD_DIM] = val_bytes
|
|
|
|
# Positive fp32 scales in [0.05, 0.55]. Byte-view into the trailing region.
|
|
scales = (
|
|
torch.rand((num_pages, PAGE_SIZE), generator=g, dtype=torch.float32).to(device)
|
|
* 0.5
|
|
+ 0.05
|
|
)
|
|
raw[:, PAGE_SIZE * HEAD_DIM :] = scales.contiguous().view(torch.uint8)
|
|
|
|
kv = raw.view(num_pages, PAGE_SIZE, 1, HEAD_DIM_WITH_SF)
|
|
return kv if dtype_view == torch.uint8 else kv.view(dtype=dtype_view)
|
|
|
|
|
|
def _build_inputs(
|
|
batch_size: int,
|
|
seq_lens: list[int],
|
|
*,
|
|
kv_dtype_view: torch.dtype,
|
|
num_heads: int = 32,
|
|
device: torch.device = torch.device("cuda"),
|
|
seed: int = 0,
|
|
):
|
|
assert len(seq_lens) == batch_size
|
|
max_seq_len = max(seq_lens)
|
|
max_pages = (max_seq_len + PAGE_SIZE - 1) // PAGE_SIZE
|
|
# One global page pool, batch picks its own page ids.
|
|
num_pages_total = batch_size * max_pages + 1
|
|
|
|
kvcache = _build_kvcache(
|
|
num_pages_total, dtype_view=kv_dtype_view, device=device, seed=seed
|
|
)
|
|
|
|
g = torch.Generator(device="cpu").manual_seed(seed + 1)
|
|
# Construct query as random bf16->fp8 to keep values inside fp8 range.
|
|
q_bf16 = torch.randn(
|
|
(batch_size, 1, num_heads, HEAD_DIM), generator=g, dtype=torch.float32
|
|
).to(device)
|
|
q_bf16 = q_bf16.clamp_(-2.0, 2.0)
|
|
q_fp8 = q_bf16.to(FP8_DTYPE)
|
|
if kv_dtype_view == torch.uint8:
|
|
# Query dtype isn't toggled — only kvcache. q always fp8.
|
|
pass
|
|
|
|
weight = (
|
|
torch.rand((batch_size, num_heads), generator=g, dtype=torch.float32).to(device)
|
|
* 0.5
|
|
)
|
|
|
|
seq_lens_t = torch.tensor(seq_lens, dtype=torch.int32, device=device)
|
|
|
|
# Each batch occupies its own slice of pages, randomized for realism.
|
|
page_table = torch.zeros((batch_size, max_pages), dtype=torch.int32, device=device)
|
|
for i in range(batch_size):
|
|
page_table[i] = torch.arange(
|
|
1 + i * max_pages, 1 + (i + 1) * max_pages, dtype=torch.int32, device=device
|
|
)
|
|
|
|
return q_fp8, kvcache, weight, seq_lens_t, page_table, max_seq_len
|
|
|
|
|
|
def _compare(
|
|
ref: torch.Tensor,
|
|
sm120: torch.Tensor,
|
|
seq_lens: torch.Tensor,
|
|
*,
|
|
atol: float = 1e-3,
|
|
rtol: float = 1e-3,
|
|
):
|
|
"""Compare reference (uninitialized beyond seq_len) vs SM120 (-inf beyond)."""
|
|
# Valid positions must match
|
|
batch_size, max_seq_len = ref.shape
|
|
for i in range(batch_size):
|
|
sl = int(seq_lens[i].item())
|
|
torch.testing.assert_close(
|
|
ref[i, :sl], sm120[i, :sl], atol=atol, rtol=rtol, equal_nan=False
|
|
)
|
|
# Invalid positions in SM120 output must be -inf
|
|
positions = torch.arange(max_seq_len, device=sm120.device)
|
|
invalid = positions.unsqueeze(0) >= seq_lens.unsqueeze(1)
|
|
assert torch.all(
|
|
torch.isinf(sm120[invalid]) & (sm120[invalid] < 0)
|
|
), "SM120 output must fill invalid positions with -inf"
|
|
|
|
|
|
class TestSM120PagedMqaLogitsTorch(CustomTestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
if not torch.cuda.is_available():
|
|
raise unittest.SkipTest("CUDA required")
|
|
cls.device = torch.device("cuda")
|
|
|
|
def _run_one(
|
|
self,
|
|
batch_size: int,
|
|
seq_lens: list[int],
|
|
kv_dtype_view: torch.dtype,
|
|
num_heads: int = 32,
|
|
):
|
|
q, kv, w, sl, pt, msl = _build_inputs(
|
|
batch_size,
|
|
seq_lens,
|
|
kv_dtype_view=kv_dtype_view,
|
|
num_heads=num_heads,
|
|
device=self.device,
|
|
)
|
|
# Reference (loopy)
|
|
ref = fp8_paged_mqa_logits_torch(
|
|
q,
|
|
kv,
|
|
w,
|
|
sl,
|
|
pt,
|
|
deep_gemm_metadata=None,
|
|
max_seq_len=msl,
|
|
clean_logits=False,
|
|
)
|
|
# SM120 vectorized
|
|
sm120 = fp8_paged_mqa_logits_torch_sm120(
|
|
q,
|
|
kv,
|
|
w,
|
|
sl,
|
|
pt,
|
|
deep_gemm_metadata=None,
|
|
max_seq_len=msl,
|
|
clean_logits=False,
|
|
)
|
|
_compare(ref, sm120, sl)
|
|
|
|
def test_equiv_fp8_view_bs1(self):
|
|
self._run_one(1, [128], FP8_DTYPE)
|
|
|
|
def test_equiv_fp8_view_bs4_uniform(self):
|
|
self._run_one(4, [128, 128, 128, 128], FP8_DTYPE)
|
|
|
|
def test_equiv_fp8_view_bs4_variable(self):
|
|
self._run_one(4, [40, 96, 200, 256], FP8_DTYPE)
|
|
|
|
def test_equiv_uint8_view_bs1(self):
|
|
"""Regression guard: KV cache viewed as raw uint8 (no FP8 dtype hint).
|
|
|
|
Historic bug (progress doc §4): some kernels treated uint8 bytes as raw
|
|
integers, producing garbled attention output. The vectorized impl must
|
|
produce the same numbers as the loopy reference regardless of the
|
|
caller's KV view dtype.
|
|
"""
|
|
self._run_one(1, [128], torch.uint8)
|
|
|
|
def test_equiv_uint8_view_bs4_variable(self):
|
|
self._run_one(4, [40, 96, 200, 256], torch.uint8)
|
|
|
|
def test_seq_lens_zero_remainder(self):
|
|
"""seq_len not aligned to page_size — last partial page must mask correctly."""
|
|
self._run_one(2, [65, 129], FP8_DTYPE) # 65 = 1 full page + 1 token
|
|
|
|
def test_seq_lens_full_pages(self):
|
|
self._run_one(2, [64, 192], FP8_DTYPE)
|
|
|
|
def test_seq_lens_2d_input_accepted(self):
|
|
"""SM120 impl squeezes seq_lens if dim>1 (matches indexer.py call site)."""
|
|
q, kv, w, sl, pt, msl = _build_inputs(
|
|
2, [64, 128], kv_dtype_view=FP8_DTYPE, device=self.device
|
|
)
|
|
sl_2d = sl.unsqueeze(-1) # (B, 1)
|
|
ref = fp8_paged_mqa_logits_torch(
|
|
q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
|
|
)
|
|
sm120 = fp8_paged_mqa_logits_torch_sm120(
|
|
q, kv, w, sl_2d, pt, None, max_seq_len=msl, clean_logits=False
|
|
)
|
|
_compare(ref, sm120, sl)
|
|
|
|
def test_cuda_graph_capture_and_replay(self):
|
|
"""No .item() / data-dependent control flow — must be CUDA-graph safe."""
|
|
batch_size = 2
|
|
seq_lens = [128, 192]
|
|
q, kv, w, sl, pt, msl = _build_inputs(
|
|
batch_size, seq_lens, kv_dtype_view=FP8_DTYPE, device=self.device
|
|
)
|
|
|
|
# Warmup outside graph
|
|
for _ in range(2):
|
|
_ = fp8_paged_mqa_logits_torch_sm120(
|
|
q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
# Pre-allocated output (graph replay reuses this buffer)
|
|
static_logits_holder = {}
|
|
|
|
graph = torch.cuda.CUDAGraph()
|
|
with torch.cuda.graph(graph):
|
|
out = fp8_paged_mqa_logits_torch_sm120(
|
|
q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
|
|
)
|
|
static_logits_holder["out"] = out
|
|
|
|
# Eager reference using the same inputs
|
|
ref = fp8_paged_mqa_logits_torch_sm120(
|
|
q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
|
|
)
|
|
|
|
graph.replay()
|
|
torch.cuda.synchronize()
|
|
torch.testing.assert_close(
|
|
static_logits_holder["out"], ref, atol=1e-5, rtol=1e-5
|
|
)
|
|
|
|
# Replay with a different seq_lens (in-place edit of the captured tensor)
|
|
sl_new = torch.tensor([64, 256], dtype=torch.int32, device=self.device)
|
|
sl.copy_(sl_new)
|
|
graph.replay()
|
|
torch.cuda.synchronize()
|
|
ref_new = fp8_paged_mqa_logits_torch_sm120(
|
|
q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
|
|
)
|
|
torch.testing.assert_close(
|
|
static_logits_holder["out"], ref_new, atol=1e-5, rtol=1e-5
|
|
)
|
|
|
|
def test_shape_assertions(self):
|
|
"""Wrong head_dim or block_size must raise."""
|
|
q, kv, w, sl, pt, msl = _build_inputs(
|
|
1, [64], kv_dtype_view=FP8_DTYPE, device=self.device
|
|
)
|
|
# head_dim != 128
|
|
bad_q = q[..., :64]
|
|
with self.assertRaises(AssertionError):
|
|
fp8_paged_mqa_logits_torch_sm120(
|
|
bad_q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=False
|
|
)
|
|
# clean_logits=True not supported
|
|
with self.assertRaises(AssertionError):
|
|
fp8_paged_mqa_logits_torch_sm120(
|
|
q, kv, w, sl, pt, None, max_seq_len=msl, clean_logits=True
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(unittest.main())
|