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:
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user