minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

feat: SM120 (Blackwell Desktop) support for DeepSeek-V4 inference (#24692)

Browse Source 
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
eeecho
2026-06-01 14:05:20 -07:00
committed by GitHub
co-authored by Claude Opus 4.6
parent dfa1af99f5
commit 524ba10eda
11 changed files with 2104 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/test_sm120_flash_mla.py
+465
View File
@@ -0,0 +1,465 @@
"""SM120 FlashMLA sparse decode unit tests.
Validates the SM120-specific FlashMLA implementation that replaces the upstream
`flash_mla` CUDA kernel (unavailable on SM120 / RTX PRO 6000):
- ``_gather_and_dequant``: byte-precise dequant of paged FP8 + BF16 + UE8M0 KV
cache. Covers the dtype-reinterpretation surface that caused the historic
uint8 garbled-output regression (see progress doc §4).
- ``_sm120_sparse_decode_fwd``: pure-PyTorch reference path.
- ``flash_mla_sparse_decode_triton``: tiled Triton kernel.
- ``_apply_attn_sink`` / ``_merge_partial_attn``: post-processing helpers.
- ``flash_mla_with_kvcache_sm120``: entry-point dispatch on
``SGLANG_SM120_TRITON_FLASHMLA`` selects torch/triton paths and both yield
matching output.
DSv4 cache layout (per page):
data section: page_size * 576 bytes = 64 tokens * (448 nope + 128 rope)
scale section: page_size * 8 bytes = 64 tokens * (7 UE8M0 scales + 1 pad)
total bytes: page_size * 584 = stride(0) of k_cache
"""
from __future__ import annotations
import unittest
from unittest import mock
import torch
from sglang.srt.layers.attention import flash_mla_sm120 as fmod
from sglang.srt.layers.attention.flash_mla_sm120 import (
_D,
_NOPE_DIM,
_NOPE_ROPE_STRIDE,
_NUM_TILES,
_ROPE_DIM,
_SCALE_STRIDE,
_TILE_SIZE,
_gather_and_dequant,
_sm120_sparse_decode_fwd,
flash_mla_with_kvcache_sm120,
)
from sglang.srt.layers.attention.flash_mla_sm120_triton import (
_apply_attn_sink,
_merge_partial_attn,
flash_mla_sparse_decode_triton,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=45, stage="base-b", runner_config="1-gpu-large")
# Per-token byte layout
_BYTES_PER_TOKEN = _NOPE_ROPE_STRIDE + _SCALE_STRIDE # 576 + 8 = 584
def _build_kvcache(
num_pages: int,
page_size: int,
*,
device: torch.device,
seed: int = 0,
):
"""Build a synthetic FP8/BF16/UE8M0 KV cache.
Returns the (num_pages, page_size, 1, bpt) FP8-viewed cache plus the raw
nope FP8 (float), rope BF16 (float), and UE8M0 scale (float) reference
tensors for verification.
"""
g = torch.Generator(device="cpu").manual_seed(seed)
bpt = _BYTES_PER_TOKEN # 584 — also satisfies k_cache.stride(0) requirement
raw = torch.zeros(num_pages, page_size, bpt, dtype=torch.uint8, device=device)
# ---- Nope FP8 region: per-token bytes [0:448] ----
# Stay in [0, 0x6F] to avoid NaN/Inf in float8_e4m3fn.
nope_bytes = torch.randint(
0, 0x70, (num_pages, page_size, _NOPE_DIM), generator=g, dtype=torch.uint8
).to(device)
raw[:, :, :_NOPE_DIM] = nope_bytes
# ---- Rope BF16 region: per-token bytes [448:576] = 64 bf16 values ----
rope_bf16_vals = (
torch.randn((num_pages, page_size, _ROPE_DIM), generator=g, dtype=torch.float32)
.clamp(-2.0, 2.0)
.to(torch.bfloat16)
.to(device)
)
# View bf16 as 2 bytes per value -> 128 bytes per token rope region
rope_as_uint8 = rope_bf16_vals.contiguous().view(
torch.uint8
) # (num_pages, page_size, 128)
raw[:, :, _NOPE_DIM : _NOPE_DIM + _ROPE_DIM * 2] = rope_as_uint8
# ---- Scale section: starts at page_size * 576 ----
# 7 UE8M0 bytes per token + 1 pad. Keep exponents in a sane range
# (UE8M0 byte 'b' decodes to 2**(b-127)), so 120..130 gives ~[1/128, 8].
scale_bytes = torch.randint(
120,
131,
(num_pages, page_size, _NUM_TILES),
generator=g,
dtype=torch.uint8,
).to(device)
# raw is (num_pages, page_size, bpt) with bpt=584; data ends at 576, scale region 576..584
# Per-token scale offset relative to scale section: token_idx * 8.
# Note: the scale region in raw memory is at the END of the page after the
# data section. Since raw shape is (num_pages, page_size, bpt), where
# bpt=584, scale lives in raw[:, :, _NOPE_ROPE_STRIDE : _NOPE_ROPE_STRIDE+8].
# BUT the gather code expects layout
# raw_pages: (num_pages, page_bytes)
# scale_section_offset = page_size * 576
# i.e. data for ALL tokens first, then scale for ALL tokens. So we need
# the flat per-page layout: [tok0_data(576), tok1_data, ..., tokN-1_data,
# tok0_scale(8), tok1_scale, ...].
# We constructed raw as (num_pages, page_size, bpt) which interleaves data
# and scales per token. Build a fresh buffer with the correct flat order.
flat = torch.zeros(num_pages, page_size * bpt, dtype=torch.uint8, device=device)
# Data: per-token 576 bytes contiguous
flat_data = raw[:, :, :_NOPE_ROPE_STRIDE].reshape(
num_pages, page_size * _NOPE_ROPE_STRIDE
)
flat[:, : page_size * _NOPE_ROPE_STRIDE] = flat_data
# Scales: per-token 8 bytes (only first 7 written)
scale_block = torch.zeros(
num_pages, page_size, _SCALE_STRIDE, dtype=torch.uint8, device=device
)
scale_block[:, :, :_NUM_TILES] = scale_bytes
flat[:, page_size * _NOPE_ROPE_STRIDE :] = scale_block.reshape(
num_pages, page_size * _SCALE_STRIDE
)
# View as (num_pages, page_size, 1, bpt) float8_e4m3fn
k_cache = flat.view(num_pages, page_size, 1, bpt).view(torch.float8_e4m3fn)
# The Triton kernel expects k_cache.stride(0) == page_size * bpt.
assert k_cache.stride(0) == page_size * bpt
# Reference dequant per token (matches what _gather_and_dequant should produce):
nope_fp8 = nope_bytes.view(
torch.float8_e4m3fn
).float() # (num_pages, page_size, 448)
scale_e8m0 = scale_bytes.view(
torch.float8_e8m0fnu
).float() # (num_pages, page_size, 7)
nope_dequant = (
nope_fp8.view(num_pages, page_size, _NUM_TILES, _TILE_SIZE)
* scale_e8m0.view(num_pages, page_size, _NUM_TILES, 1)
).view(
num_pages, page_size, _NOPE_DIM
) # float32
ref_per_token = torch.cat(
[nope_dequant.to(torch.bfloat16), rope_bf16_vals], dim=-1
) # (num_pages, page_size, 512) bf16
return k_cache, ref_per_token
def _build_q_indices(
batch_size: int,
num_heads: int,
topk: int,
num_pages: int,
page_size: int,
*,
device: torch.device,
seed: int = 1,
):
g = torch.Generator(device="cpu").manual_seed(seed)
q = (
torch.randn((batch_size, 1, num_heads, _D), generator=g, dtype=torch.float32)
.clamp(-1.5, 1.5)
.to(torch.bfloat16)
.to(device)
)
# Each batch picks `topk` random valid token-level indices into the pool
pool_size = num_pages * page_size
indices = torch.zeros((batch_size, 1, topk), dtype=torch.int32, device=device)
for b in range(batch_size):
perm = torch.randperm(pool_size, generator=g)[:topk]
indices[b, 0] = perm.to(device=device, dtype=torch.int32)
return q, indices
class TestGatherAndDequant(CustomTestCase):
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA required")
cls.device = torch.device("cuda")
def test_basic_dequant_matches_manual(self):
"""Per-token output equals nope_fp8 * scale concatenated with rope_bf16."""
num_pages, page_size = 3, 64
k_cache, ref_per_token = _build_kvcache(
num_pages, page_size, device=self.device, seed=42
)
# Pick 1 index per page (positions 0, 32, 63)
token_ids = torch.tensor(
[[0, 1 * page_size + 32, 2 * page_size + 63]],
dtype=torch.int32,
device=self.device,
)
out = _gather_and_dequant(k_cache, token_ids, page_size)
self.assertEqual(out.shape, (1, 3, _D))
# Expected entries
expected = torch.stack(
[
ref_per_token[0, 0],
ref_per_token[1, 32],
ref_per_token[2, 63],
],
dim=0,
).unsqueeze(0)
# bf16 dequant: allow up to 1 ULP per element. Scales are integral
# powers of 2 so the only loss is the fp8 -> bf16 mantissa rounding.
torch.testing.assert_close(out, expected, atol=1e-2, rtol=1e-2)
def test_dequant_handles_full_page_range(self):
"""All token positions in a single page produce correct output."""
num_pages, page_size = 2, 64
k_cache, ref_per_token = _build_kvcache(
num_pages, page_size, device=self.device, seed=7
)
token_ids = torch.arange(page_size, dtype=torch.int32, device=self.device).view(
1, page_size
)
out = _gather_and_dequant(k_cache, token_ids, page_size)
torch.testing.assert_close(
out, ref_per_token[0].unsqueeze(0), atol=1e-2, rtol=1e-2
)
class TestSparseDecodeTritonVsTorch(CustomTestCase):
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA required")
cls.device = torch.device("cuda")
def _run(
self,
batch_size: int = 2,
num_heads: int = 4,
topk: int = 64,
page_size: int = 64,
num_pages: int = 4,
seed: int = 11,
with_topk_length: bool = True,
with_attn_sink: bool = False,
):
k_cache, _ = _build_kvcache(num_pages, page_size, device=self.device, seed=seed)
q, indices = _build_q_indices(
batch_size,
num_heads,
topk,
num_pages,
page_size,
device=self.device,
seed=seed + 100,
)
topk_length = None
if with_topk_length:
topk_length = torch.tensor(
[topk // 2, topk] if batch_size == 2 else [topk] * batch_size,
dtype=torch.int32,
device=self.device,
)
attn_sink = None
if with_attn_sink:
attn_sink = torch.full(
(num_heads,), -1.5, dtype=torch.float32, device=self.device
)
softmax_scale = _D ** (-0.5)
# Production passes config.v_head_dim == _D (512). Both impls return
# (B, 1, H, _D); the Triton kernel always computes the full nope+rope.
head_dim_v = _D
# PyTorch reference
ref_out, ref_lse = _sm120_sparse_decode_fwd(
q,
k_cache,
indices,
topk_length,
attn_sink,
head_dim_v=head_dim_v,
softmax_scale=softmax_scale,
)
# Triton
tri_out, tri_lse = flash_mla_sparse_decode_triton(
q,
k_cache,
indices,
topk_length,
attn_sink,
head_dim_v=head_dim_v,
softmax_scale=softmax_scale,
)
# Outputs both bf16, allow drift from online softmax base-2 vs base-e
# math + bf16 mantissa precision (~7 bits). 5e-2 abs/rel is generous
# but accounts for accumulated rounding across topk tokens.
torch.testing.assert_close(
tri_out.to(torch.float32),
ref_out.to(torch.float32),
atol=5e-2,
rtol=5e-2,
)
return ref_out, tri_out, ref_lse, tri_lse
def test_triton_vs_torch_basic(self):
self._run()
def test_triton_vs_torch_no_topk_length(self):
self._run(with_topk_length=False)
def test_triton_vs_torch_with_attn_sink(self):
self._run(with_attn_sink=True)
def test_triton_vs_torch_small_topk(self):
self._run(topk=32, num_heads=2)
def test_triton_vs_torch_negative_indices_are_masked(self):
"""Indices < 0 are 'invalid' and must contribute zero to output."""
k_cache, _ = _build_kvcache(4, 64, device=self.device, seed=3)
q, indices = _build_q_indices(2, 4, 32, 4, 64, device=self.device, seed=99)
# Half of each batch's indices set to -1
indices[:, :, 16:] = -1
topk_length = torch.tensor([32, 32], dtype=torch.int32, device=self.device)
ref_out, _ = _sm120_sparse_decode_fwd(
q, k_cache, indices, topk_length, None, _D, _D**-0.5
)
tri_out, _ = flash_mla_sparse_decode_triton(
q, k_cache, indices, topk_length, None, _D, _D**-0.5
)
torch.testing.assert_close(
tri_out.to(torch.float32),
ref_out.to(torch.float32),
atol=5e-2,
rtol=5e-2,
)
class TestApplyAttnSink(CustomTestCase):
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA required")
cls.device = torch.device("cuda")
def test_sink_zero_means_full_dampen_to_half(self):
"""attn_sink == lse implies sink contributes equal weight: output halves."""
B, H, D = 1, 2, 8
lse = torch.zeros(B, 1, H, device=self.device, dtype=torch.float32)
attn_sink = torch.zeros(H, device=self.device, dtype=torch.float32)
out = torch.ones(B, 1, H, D, device=self.device, dtype=torch.bfloat16)
new_out, new_lse = _apply_attn_sink(out, lse, attn_sink)
# combined_lse = logaddexp(0, 0) = log(2). w = exp(-log2) = 0.5.
torch.testing.assert_close(
new_out.float(),
torch.full_like(new_out.float(), 0.5),
atol=2e-3,
rtol=2e-3,
)
torch.testing.assert_close(
new_lse,
torch.full_like(new_lse, 0.6931472), # log(2)
atol=1e-5,
rtol=1e-5,
)
def test_sink_dead_lse_stays_zero(self):
"""lse == -inf (no valid tokens) -> output stays zero (weight=0)."""
B, H, D = 1, 2, 4
lse = torch.full(
(B, 1, H), float("-inf"), device=self.device, dtype=torch.float32
)
attn_sink = torch.zeros(H, device=self.device, dtype=torch.float32)
out = torch.ones(B, 1, H, D, device=self.device, dtype=torch.bfloat16)
new_out, _ = _apply_attn_sink(out, lse, attn_sink)
torch.testing.assert_close(
new_out.float(), torch.zeros_like(new_out.float()), atol=0, rtol=0
)
class TestMergePartialAttn(CustomTestCase):
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA required")
cls.device = torch.device("cuda")
def test_equal_lse_arithmetic_mean(self):
"""lse1 == lse2 -> merged output is the arithmetic mean of out1, out2."""
B, H, D = 1, 2, 4
lse = torch.zeros(B, 1, H, device=self.device, dtype=torch.float32)
out1 = torch.ones(B, 1, H, D, device=self.device, dtype=torch.bfloat16)
out2 = torch.full((B, 1, H, D), 3.0, device=self.device, dtype=torch.bfloat16)
merged, merged_lse = _merge_partial_attn(out1, lse, out2, lse)
torch.testing.assert_close(
merged.float(),
torch.full_like(merged.float(), 2.0),
atol=1e-2,
rtol=1e-2,
)
torch.testing.assert_close(
merged_lse,
torch.full_like(merged_lse, 0.6931472), # log(2)
atol=1e-5,
rtol=1e-5,
)
def test_one_dead_branch_passes_through(self):
"""If lse2 == -inf, merged equals out1."""
B, H, D = 1, 1, 4
lse1 = torch.zeros(B, 1, H, device=self.device, dtype=torch.float32)
lse2 = torch.full(
(B, 1, H), float("-inf"), device=self.device, dtype=torch.float32
)
out1 = torch.full((B, 1, H, D), 2.5, device=self.device, dtype=torch.bfloat16)
out2 = torch.full((B, 1, H, D), 99.0, device=self.device, dtype=torch.bfloat16)
merged, merged_lse = _merge_partial_attn(out1, lse1, out2, lse2)
torch.testing.assert_close(merged.float(), out1.float(), atol=1e-3, rtol=1e-3)
torch.testing.assert_close(merged_lse, lse1, atol=1e-5, rtol=1e-5)
class TestEntryPointDispatch(CustomTestCase):
@classmethod
def setUpClass(cls):
if not torch.cuda.is_available():
raise unittest.SkipTest("CUDA required")
cls.device = torch.device("cuda")
def test_torch_backend_matches_triton_backend(self):
"""SGLANG_SM120_TRITON_FLASHMLA toggles backend; both return matching out."""
k_cache, _ = _build_kvcache(4, 64, device=self.device, seed=5)
q, indices = _build_q_indices(1, 4, 32, 4, 64, device=self.device, seed=13)
topk_length = torch.tensor([32], dtype=torch.int32, device=self.device)
kwargs = dict(
q=q,
k_cache=k_cache,
indices=indices,
topk_length=topk_length,
attn_sink=None,
head_dim_v=_D,
softmax_scale=_D**-0.5,
)
with mock.patch.object(fmod, "_sm120_default_backend", "torch"):
out_torch, _ = flash_mla_with_kvcache_sm120(**kwargs)
with mock.patch.object(fmod, "_sm120_default_backend", "triton"):
out_triton, _ = flash_mla_with_kvcache_sm120(**kwargs)
torch.testing.assert_close(
out_torch.to(torch.float32),
out_triton.to(torch.float32),
atol=5e-2,
rtol=5e-2,
)
if __name__ == "__main__":
import sys
sys.exit(unittest.main())
test/registered/kernels/test_sm120_paged_mqa_logits.py
+314
View File
@@ -0,0 +1,314 @@
"""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())