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

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:
2026-06-12 04:02:19 +00:00
parent 9c6cf278b0
commit d66758d929
3 changed files with 29 additions and 10 deletions

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)

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,