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

Run the symm staging fill on byte views (no fp8 index_copy kernel)

Browse Source 
First e2e launch crashed every CP rank at the symm token fill:
index_copy_cuda is not implemented for Float8_e4m3fn — the production
KV pool dtype, which the 8-rank test missed by building its pools as
uint8.  The fill is a whole-token-row copy, so it is dtype-agnostic:
both the staging span and the current rows now go through uint8 views.

The 8-rank test now builds the KV pools and current rows as
float8_e4m3fn (payloads constructed as bytes, compared as bytes) so
the production dtype is what every phase exercises.

Validated on g0033: all four 8-rank byte-exactness phases under fp8.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
leavelet
2026-06-12 04:02:19 +00:00
co-authored by Claude Fable 5
parent 9c6cf278b0
commit d66758d929
3 changed files with 29 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
+7 -3
View File
@@ -1001,14 +1001,18 @@ class CpSharedKVMlaPrefetcher:
)
num_rows = int(state.staging_current_rows.numel())
if num_rows > 0:
staging_rows = staging_span.view(kv_cache.dtype).reshape(
# Byte view on both sides: index_copy_ has no fp8 CUDA
# kernel, and the copy is dtype-agnostic (whole token rows).
staging_rows = staging_span.view(
(state.num_current_pages + 1) * self.page_size,
*kv_cache.shape[1:],
state.page_nbytes // self.page_size,
)
staging_rows.index_copy_(
0,
state.staging_current_rows,
current_kv_cache[:num_rows],
current_kv_cache[:num_rows]
.reshape(num_rows, -1)
.view(torch.uint8),
)
cp_symm_barrier(
staging.flag_ptrs, self_rank=int(self.layout.cp_rank)
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
+7 -3
View File
@@ -4876,14 +4876,18 @@ def _compose_token_kv_partial_current_v2(
)
num_rows = int(fill_state.staging_current_rows.numel())
if num_rows > 0:
staging_rows = staging_span.view(kv_cache.dtype).reshape(
# Byte view on both sides: index_copy_ has no fp8 CUDA kernel,
# and the copy is dtype-agnostic anyway (whole token rows).
staging_rows = staging_span.view(
(int(plan.num_current_pages) + 1) * page_size,
*kv_cache.shape[1:],
kv_page_nbytes // page_size,
)
staging_rows.index_copy_(
0,
fill_state.staging_current_rows,
current_kv_cache[:num_rows],
current_kv_cache[:num_rows]
.reshape(num_rows, -1)
.view(torch.uint8),
)
_symm_barrier_and_gather_all(
kernels=kernels,
test/manual/test_cp_shared_kv_compose_v2_8rank.py
+15 -4
View File
@@ -110,6 +110,8 @@ def _build_scenario(rank: int, cp_size: int, device: torch.device):
logical_locs = logical_locs.to(device)
physical_pages = (next_page // cp_size + 3)
# Production KV dtype: float8_e4m3fn (caught index_copy_'s missing fp8
# kernel) — build payloads as bytes, view as fp8.
kv_cache = torch.zeros(
(physical_pages * page_size, 1, kv_dim), dtype=torch.uint8, device=device
)
@@ -154,7 +156,9 @@ def _build_scenario(rank: int, cp_size: int, device: torch.device):
.remainder_(249)
.to(torch.uint8)
.to(device)
.view(torch.float8_e4m3fn)
)
kv_cache = kv_cache.view(torch.float8_e4m3fn)
prefix_slot_spans = runtime.build_batch_prefix_slot_spans(
logical_pages=logical_pages,
@@ -347,7 +351,9 @@ def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None:
.remainder_(247)
.to(torch.uint8)
.to(device)
.view(torch.float8_e4m3fn)
)
kv_cache = kv_cache.view(torch.float8_e4m3fn)
current_k = (
(
torch.arange(rows, dtype=torch.int64).view(-1, 1) * 23
@@ -435,7 +441,7 @@ def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None:
end_slot=prefix_pages,
)
prefix_rows = rt.slot_range_to_token_slice(page_size, 0, prefix_pages)
dist.all_reduce(dense_kv[prefix_rows])
dist.all_reduce(dense_kv[prefix_rows].view(torch.uint8))
mla_event = torch.cuda.Event()
mla_event.record()
mla.handles[61] = CpSharedKVMlaPrefetchHandle(
@@ -508,8 +514,10 @@ def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None:
raise AssertionError(
f"rank{rank}: prefetch index symm mismatch at pages {bad.cpu().tolist()}"
)
if not torch.equal(ref_kv, sym_kv):
diff = (ref_kv != sym_kv).any(dim=-1).any(dim=-1)
ref_kv_b = ref_kv.view(torch.uint8)
sym_kv_b = sym_kv.view(torch.uint8)
if not torch.equal(ref_kv_b, sym_kv_b):
diff = (ref_kv_b != sym_kv_b).any(dim=-1).any(dim=-1)
bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
raise AssertionError(
f"rank{rank}: prefetch mla symm mismatch at rows {bad} "
@@ -562,6 +570,8 @@ def main() -> None:
assert torch.equal(ref_locs, v2_locs), (
f"rank{rank}: locs mismatch legacy vs v2"
)
ref_kv = ref_kv.view(torch.uint8)
v2_kv = v2_kv.view(torch.uint8)
if not torch.equal(ref_kv, v2_kv):
diff = (ref_kv != v2_kv).any(dim=-1).any(dim=-1)
bad_rows = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
@@ -571,7 +581,7 @@ def main() -> None:
)
# Cross-rank: every rank must hold the SAME composed buffer.
ref_sum = ref_kv.to(torch.float64).sum()
ref_sum = ref_kv.view(torch.uint8).to(torch.float64).sum()
sums = torch.zeros(world, dtype=torch.float64, device=device)
dist.all_gather_into_tensor(
sums, ref_sum.reshape(1).to(device)
@@ -596,6 +606,7 @@ def main() -> None:
symm_kv, symm_locs = _compose(
s, layer_id=layer_id, writers=s["current_page_writer_ranks"]
)
symm_kv = symm_kv.view(torch.uint8)
torch.cuda.synchronize()
assert torch.equal(ref_locs, symm_locs), (
f"rank{rank} layer{layer_id} [{tag}]: symm locs mismatch"