[CP] 1/N: Support MLA Prefill Context Parallel (#23292)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
81cd338fcc
commit
b0ce16d0c5
@@ -0,0 +1,143 @@
|
||||
"""
|
||||
FA3 parity test for `prepare_context_parallel_metadata`.
|
||||
|
||||
Drives the real function and feeds its `kv_len_prev/next_tensor` into FA3
|
||||
via `flash_attn_with_kvcache`. Compares per-rank CP output against a
|
||||
full-sequence FA3 reference computed over the unpadded `(prefix + extend)`
|
||||
KV. Any discrepancy indicates the metadata function emitted wrong
|
||||
`cache_seqlens` for at least one rank.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.utils.cp_utils import prepare_context_parallel_metadata
|
||||
from sglang.srt.utils.common import ceil_align
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=5, stage="extra-a", runner_config="1-gpu-large")
|
||||
|
||||
_DSA_UTILS = "sglang.srt.layers.attention.dsa.utils"
|
||||
_DEVICE = "cuda"
|
||||
_DTYPE = torch.bfloat16
|
||||
_HEAD_NUM = 8
|
||||
_HEAD_DIM = 128
|
||||
_SCALE = _HEAD_DIM**-0.5
|
||||
|
||||
|
||||
class TestCPPrefixLenFA3Parity(CustomTestCase):
|
||||
"""Per-rank FA3 output under CP must match a full-sequence reference."""
|
||||
|
||||
def _run_parity(self, prefix_len: int, extend_len: int, cp_size: int):
|
||||
from sgl_kernel.flash_attn import flash_attn_with_kvcache
|
||||
|
||||
torch.manual_seed(extend_len * 1_000_003 + prefix_len * 101 + cp_size)
|
||||
|
||||
padded_extend = ceil_align(extend_len, cp_size)
|
||||
pad = padded_extend - extend_len
|
||||
self.assertGreaterEqual(
|
||||
padded_extend,
|
||||
2 * cp_size,
|
||||
"runtime `can_cp_split` would skip this case; pick a larger extend",
|
||||
)
|
||||
|
||||
# Reference: one full-sequence FA3 call over the unpadded KV.
|
||||
q_full = torch.randn(
|
||||
extend_len, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
|
||||
)
|
||||
k_full = torch.randn(
|
||||
prefix_len + extend_len, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
|
||||
)
|
||||
v_full = torch.randn(
|
||||
prefix_len + extend_len, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
|
||||
)
|
||||
ref = flash_attn_with_kvcache(
|
||||
q=q_full.unsqueeze(0),
|
||||
k_cache=k_full.unsqueeze(0),
|
||||
v_cache=v_full.unsqueeze(0),
|
||||
cache_seqlens=torch.tensor(
|
||||
[k_full.shape[0]], dtype=torch.int32, device=_DEVICE
|
||||
),
|
||||
softmax_scale=_SCALE,
|
||||
causal=True,
|
||||
).squeeze(0)
|
||||
|
||||
# CP path sees tensors padded to `ceil_align(extend, cp_size)`,
|
||||
# matching what `prepare_mlp_sync_batch` does in production.
|
||||
zeros = torch.zeros(pad, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE)
|
||||
q_padded = torch.cat([q_full, zeros], dim=0)
|
||||
k_padded = torch.cat([k_full, zeros], dim=0)
|
||||
v_padded = torch.cat([v_full, zeros], dim=0)
|
||||
|
||||
seqs_len = [prefix_len + extend_len]
|
||||
extend_lens = [extend_len]
|
||||
|
||||
def _call_meta(rank: int):
|
||||
return prepare_context_parallel_metadata(
|
||||
padded_extend, rank, cp_size, seqs_len, extend_lens=extend_lens
|
||||
)
|
||||
|
||||
# Exercise the non-DSA branch; the DSA branch uses a separate
|
||||
# `prefix_len` pathway re-added by `_get_topk_ragged_with_cp`.
|
||||
with (
|
||||
patch(f"{_DSA_UTILS}.is_dsa_enable_prefill_cp", return_value=False),
|
||||
patch(
|
||||
f"{_DSA_UTILS}.is_dsa_prefill_cp_round_robin_split",
|
||||
return_value=False,
|
||||
),
|
||||
):
|
||||
meta0 = _call_meta(0)
|
||||
cp_segment_num = 2 * cp_size
|
||||
blocks_q = list(torch.split(q_padded, meta0.split_list, dim=0))
|
||||
outs = [None] * cp_segment_num
|
||||
|
||||
for rank in range(cp_size):
|
||||
meta = meta0 if rank == 0 else _call_meta(rank)
|
||||
for idx, cs_tensor in (
|
||||
(rank, meta.kv_len_prev_tensor),
|
||||
(cp_size * 2 - rank - 1, meta.kv_len_next_tensor),
|
||||
):
|
||||
if meta0.split_list[idx] == 0:
|
||||
outs[idx] = torch.empty(
|
||||
0, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
|
||||
)
|
||||
continue
|
||||
outs[idx] = flash_attn_with_kvcache(
|
||||
q=blocks_q[idx].unsqueeze(0),
|
||||
k_cache=k_padded.unsqueeze(0),
|
||||
v_cache=v_padded.unsqueeze(0),
|
||||
cache_seqlens=cs_tensor,
|
||||
softmax_scale=_SCALE,
|
||||
causal=True,
|
||||
).squeeze(0)
|
||||
|
||||
cp_out = torch.cat(outs, dim=0)
|
||||
err = (cp_out[:extend_len].float() - ref.float()).abs().max().item()
|
||||
|
||||
self.assertLess(
|
||||
err,
|
||||
1e-2,
|
||||
f"CP output diverges from full-sequence FA3 reference by "
|
||||
f"max_err={err:.5f} "
|
||||
f"(prefix_len={prefix_len}, extend_len={extend_len}, "
|
||||
f"cp_size={cp_size}, pad={pad})",
|
||||
)
|
||||
|
||||
def test_cp2_prefix1_extend3(self):
|
||||
"""cp_size=2, prefix_len=1, extend_len=3 (pad=1)."""
|
||||
self._run_parity(prefix_len=1, extend_len=3, cp_size=2)
|
||||
|
||||
def test_cp4_prefix1_extend7(self):
|
||||
"""cp_size=4, prefix_len=1, extend_len=7 (pad=1)."""
|
||||
self._run_parity(prefix_len=1, extend_len=7, cp_size=4)
|
||||
|
||||
def test_cp8_prefix1_extend17(self):
|
||||
"""cp_size=8, prefix_len=1, extend_len=17 (pad=7)."""
|
||||
self._run_parity(prefix_len=1, extend_len=17, cp_size=8)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,203 @@
|
||||
"""FA3 numerical parity for MLA prefill CP.
|
||||
|
||||
Verifies the rank-local zigzag-split FA3 path (``_mla_cp_attn`` +
|
||||
``cp_attn_forward_extend`` in ``flashattention_backend.py``) matches a
|
||||
single non-CP ``flash_attn_with_kvcache`` over the full sequence.
|
||||
|
||||
Single-process, single-layer, pre-populated paged KV cache. Requires
|
||||
FA3 ver=3 (Hopper+).
|
||||
"""
|
||||
|
||||
import math
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.utils.cp_utils import (
|
||||
ContextParallelMetadata,
|
||||
cp_attn_forward_extend,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-large")
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip(reason="CUDA required for FA3", allow_module_level=True)
|
||||
|
||||
_cap = torch.cuda.get_device_capability(0)
|
||||
if _cap[0] < 9:
|
||||
pytest.skip(
|
||||
reason=f"FA3 ver=3 requires Hopper (sm90+); got sm{_cap[0]}{_cap[1]}",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
try:
|
||||
from sgl_kernel.flash_attn import flash_attn_with_kvcache
|
||||
except ImportError as e:
|
||||
pytest.skip(
|
||||
reason=f"sgl_kernel.flash_attn unavailable: {e}",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
DEVICE = torch.device("cuda")
|
||||
DTYPE = torch.bfloat16
|
||||
|
||||
# Default shape is DeepSeek V3/R1 TP=8 MLA: 16 heads, v=512, rope=64.
|
||||
NUM_HEADS = 16
|
||||
V_HEAD_DIM = 512
|
||||
QK_ROPE_HEAD_DIM = 64
|
||||
PAGE_SIZE = 1
|
||||
|
||||
|
||||
def _build_cache_and_q(seq_len):
|
||||
"""Pre-populated paged KV cache + full-sequence q.
|
||||
|
||||
Pre-population mirrors upstream ``rebuild_cp_kv_cache``, which all-gathers
|
||||
rank-local KV into the global pool before the attention call, so each
|
||||
rank's FA3 invocation sees the same fully-populated cache.
|
||||
"""
|
||||
num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
|
||||
c_kv_cache = torch.randn(
|
||||
num_pages, PAGE_SIZE, 1, V_HEAD_DIM, dtype=DTYPE, device=DEVICE
|
||||
)
|
||||
k_rope_cache = torch.randn(
|
||||
num_pages, PAGE_SIZE, 1, QK_ROPE_HEAD_DIM, dtype=DTYPE, device=DEVICE
|
||||
)
|
||||
q_nope = torch.randn(seq_len, NUM_HEADS, V_HEAD_DIM, dtype=DTYPE, device=DEVICE)
|
||||
q_rope = torch.randn(
|
||||
seq_len, NUM_HEADS, QK_ROPE_HEAD_DIM, dtype=DTYPE, device=DEVICE
|
||||
)
|
||||
page_table = torch.arange(num_pages, dtype=torch.int32, device=DEVICE).unsqueeze(0)
|
||||
return c_kv_cache, k_rope_cache, q_nope, q_rope, page_table
|
||||
|
||||
|
||||
def _full_seq_attn(
|
||||
seq_len, q_nope, q_rope, c_kv_cache, k_rope_cache, page_table, softmax_scale
|
||||
):
|
||||
"""Non-CP reference: single flash_attn_with_kvcache over the full seq."""
|
||||
return flash_attn_with_kvcache(
|
||||
q=q_rope,
|
||||
qv=q_nope,
|
||||
k_cache=k_rope_cache,
|
||||
v_cache=c_kv_cache,
|
||||
page_table=page_table,
|
||||
cache_seqlens=torch.tensor([seq_len], dtype=torch.int32, device=DEVICE),
|
||||
cu_seqlens_q=torch.tensor([0, seq_len], dtype=torch.int32, device=DEVICE),
|
||||
cu_seqlens_k_new=None,
|
||||
max_seqlen_q=seq_len,
|
||||
softmax_scale=softmax_scale,
|
||||
causal=True,
|
||||
ver=3,
|
||||
)
|
||||
|
||||
|
||||
def _cp_attn_for_rank(
|
||||
rank,
|
||||
cp_size,
|
||||
block_size,
|
||||
q_nope,
|
||||
q_rope,
|
||||
c_kv_cache,
|
||||
k_rope_cache,
|
||||
page_table,
|
||||
softmax_scale,
|
||||
):
|
||||
"""Run the rank-local CP closure from ``flashattention_backend.py``.
|
||||
|
||||
Zigzag layout: rank r gets blocks [r, num_blocks - 1 - r] where
|
||||
num_blocks = cp_size * 2. kv_len for each half is the cumulative KV
|
||||
extent through the end of that block.
|
||||
"""
|
||||
num_blocks = cp_size * 2
|
||||
b_prev, b_next = rank, num_blocks - 1 - rank
|
||||
prev_slice = slice(b_prev * block_size, (b_prev + 1) * block_size)
|
||||
next_slice = slice(b_next * block_size, (b_next + 1) * block_size)
|
||||
|
||||
q_nope_local = torch.cat([q_nope[prev_slice], q_nope[next_slice]], dim=0)
|
||||
q_rope_local = torch.cat([q_rope[prev_slice], q_rope[next_slice]], dim=0)
|
||||
q_fused = torch.cat([q_nope_local, q_rope_local], dim=-1)
|
||||
|
||||
cp_meta = ContextParallelMetadata(
|
||||
kv_len_prev_tensor=torch.tensor(
|
||||
[(b_prev + 1) * block_size], dtype=torch.int32, device=DEVICE
|
||||
),
|
||||
kv_len_next_tensor=torch.tensor(
|
||||
[(b_next + 1) * block_size], dtype=torch.int32, device=DEVICE
|
||||
),
|
||||
actual_seq_q_prev=block_size,
|
||||
actual_seq_q_next=block_size,
|
||||
)
|
||||
fb = SimpleNamespace(attn_cp_metadata=cp_meta)
|
||||
|
||||
def _mla_cp_attn(q_chunk, cu_seqlens_q_cp, cache_seqlens_cp, max_seqlen_q_cp):
|
||||
q_nope_chunk = q_chunk[..., :V_HEAD_DIM]
|
||||
q_rope_chunk = q_chunk[..., V_HEAD_DIM:]
|
||||
return flash_attn_with_kvcache(
|
||||
q=q_rope_chunk,
|
||||
qv=q_nope_chunk,
|
||||
k_cache=k_rope_cache,
|
||||
v_cache=c_kv_cache,
|
||||
page_table=page_table,
|
||||
cache_seqlens=cache_seqlens_cp,
|
||||
cu_seqlens_q=cu_seqlens_q_cp,
|
||||
cu_seqlens_k_new=None,
|
||||
max_seqlen_q=max_seqlen_q_cp,
|
||||
softmax_scale=softmax_scale,
|
||||
causal=True,
|
||||
ver=3,
|
||||
)
|
||||
|
||||
local_out = cp_attn_forward_extend(fb, q_fused, DEVICE, _mla_cp_attn)
|
||||
return local_out, prev_slice, next_slice
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cp_size, block_size",
|
||||
[
|
||||
(2, 64), # DSv3 TP=8 baseline
|
||||
(2, 128), # longer per-block seq
|
||||
(4, 32), # multi-rank zigzag: rank r gets blocks [r, 7-r]
|
||||
],
|
||||
)
|
||||
def test_cp_parity(cp_size, block_size):
|
||||
torch.manual_seed(0)
|
||||
seq_len = block_size * cp_size * 2
|
||||
softmax_scale = 1.0 / math.sqrt(V_HEAD_DIM + QK_ROPE_HEAD_DIM)
|
||||
|
||||
c_kv_cache, k_rope_cache, q_nope, q_rope, page_table = _build_cache_and_q(seq_len)
|
||||
ref_out = _full_seq_attn(
|
||||
seq_len, q_nope, q_rope, c_kv_cache, k_rope_cache, page_table, softmax_scale
|
||||
)
|
||||
|
||||
for rank in range(cp_size):
|
||||
local_out, prev_slice, next_slice = _cp_attn_for_rank(
|
||||
rank,
|
||||
cp_size,
|
||||
block_size,
|
||||
q_nope,
|
||||
q_rope,
|
||||
c_kv_cache,
|
||||
k_rope_cache,
|
||||
page_table,
|
||||
softmax_scale,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
local_out[:block_size],
|
||||
ref_out[prev_slice],
|
||||
rtol=1e-3,
|
||||
atol=5e-3,
|
||||
msg=f"rank={rank} prev-half mismatch",
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
local_out[block_size:],
|
||||
ref_out[next_slice],
|
||||
rtol=1e-3,
|
||||
atol=5e-3,
|
||||
msg=f"rank={rank} next-half mismatch",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
Reference in New Issue
Block a user