diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index 51207448a..a3f5a3587 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -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) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 67c888939..db7c39c5a 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -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, diff --git a/test/manual/test_cp_shared_kv_compose_v2_8rank.py b/test/manual/test_cp_shared_kv_compose_v2_8rank.py index afea40bd5..47ea1cda1 100644 --- a/test/manual/test_cp_shared_kv_compose_v2_8rank.py +++ b/test/manual/test_cp_shared_kv_compose_v2_8rank.py @@ -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"