First captured graph keeps a transient cuBLAS split-K workspace: Xid 31 on DeepSeek V4.1 TP4 with --prefill-chunk-size 8192 at 7f534bb6

Aperta
  • Qualcosa non va?
#266 3 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
43/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
python, pytorch, rust
Ambito
backend, performance

Direzione di ricerca

Start with ignition/preparation/capture.py and trace capture_forwards, warmup(), and how the first cuBLAS workspace is allocated and recorded during graph capture. The report includes a proposed workspace clearing change, but its export and runtime checks are still pending. Done means verifying that all graph workspace pointers are backed by recorded allocations and that the 8192 cold-prefill case runs across ranks.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

Hi, and thank you for Hyperflare. We run DeepSeek-V4.1-Flash-1017 at TP=4 on four DGX Spark (GB10) nodes on the RoCE ring. At 7f534bb6, a package exported with --prefill-chunk-size 8192 crashes all four ranks on the first prompt that selects the prefill8192 graph. We traced the cause to the cuBLAS workspace of the first captured graph. The defect also affects the default chunk size, but there it is silent.

Symptom

  • A cold prompt of 7,014 tokens works (the prefill7168 graph). A cold prompt of 7,900 tokens (prefill8192) fails every time, on every rank: 
    NVRM: Xid (PCI:000f:01:00): 31, pid=..., name=flare, intr 00000000. MMU Fault: ENGINE GRAPHICS GPC0 GPCCLIENT_T1_0 faulted @ 0x812_d3a00000. Fault is of type FAULT_PDE ACCESS_TYPE_VIRT_WRITE
Error: "native replay did not drain: cuEventQuery failed (700, CUDA_ERROR_ILLEGAL_ADDRESS): an illegal memory access was encountered"
  • The fault address is the same on every run. It does not depend on the KV cache: the fault stays with --kv-cache-limit context (1 GiB of KV).
  • The same 8192 chunk worked on 085ae4a7. A chunk-4096 export of 7f534bb6 serves normally.

Where the address comes from

flare inspect --memory-map shows that no recorded allocation covers 0x812d3a00000. It lies in a 104 MiB gap between two recorded device allocations. We decoded rank 0's graphs.safetensors (zstd, then length-prefixed Node records) and searched the kernel argument bytes. Exactly 10 arguments hold this address, all in prefill8192, in 5 pairs:

Function Argument
nvjet_sm121_tst_mma_32x64x64_8_32x16x64_tmaAB_alignCD4_splitK_TNNN 384-byte parameter, offset 256
cublasLt::splitKreduce_kernel<32,16,int,float,bf16,...> argument 1 (the split-K workspace)

The same split-K pair occurs in every graph of the package (prefill*, decode8, speculative*, packed-speculative*). The workspace argument has two values:

Graphs Workspace Owner
prefill8192 (the first captured graph) 0x812d3a00000 none (unmapped at serve time)
every other graph 0x812d8a00000 device-570, recorded

The chunk-4096 export of the same pin has the same pattern. There, prefill4096 (captured first) uses 0x8126cc00000, the start of device-547. Every other graph uses device-548. So in that package, the first graph's split-K partials go to storage that another recorded allocation owns. It does not crash, but it aliases.

Cause (our reading of ignition/preparation/capture.py)

capture_forwards runs every forward once in "Warm serving graphs" inside warmup(), which is warmup_storage.transient. PyTorch caches one cuBLAS workspace per (handle, stream), and the first cuBLAS call on the capture stream creates it there, in transient storage. The first capture reuses that cached workspace, so its address goes into the graph. After that capture, PyTorch's cached workspace is gone, and every later capture allocates its workspace from the graph pool, which export records. Export then releases transient storage, so the first graph replays a pointer that no allocation owns. The 8192 package faults because nothing covers the address. The 4096 package survives because a later allocation happens to cover it.

The 085ae4a7 capture code also warmed inside warmup_storage.transient. So we think the defect is older, and only the memory layout changed whether it shows.

Change we are testing

Clear the cached workspaces once, after the warmups and before the first capture:

--- a/ignition/preparation/capture.py
+++ b/ignition/preparation/capture.py
@@ capture_forwards
             stream.synchronize()
             stage.update(advance=1)
+    # PyTorch caches one cuBLAS workspace per handle and stream, and the eager warmups above
+    # created it inside transient storage. The first capture would otherwise bake that address
+    # into its split-K GEMMs and export it unrecorded; later captures already take theirs from
+    # the graph pool. Drop the cached workspaces so every capture allocates its own there.
+    torch.cuda.synchronize(device)
+    torch._C._cuda_clearCublasWorkspaces()
     with phase(progress, label, total=len(names)) as stage:

A chunk-8192 export with this change is running now. We will add the result here: the pointer check on all ranks, a cold-prefill ladder above 8K tokens, and long-prefill timing. If you prefer another fix, for example creating the workspace in persistent storage before the warmups or recording it as an allocation, we can test that instead.

A check that may help export validation

We wrote a small script that walks every kernel argument of a package. It reports 8-byte values that fall inside the rank's reservations but outside every recorded allocation. On the faulting package it reports exactly the 10 arguments above. Its only other hits are constants in the W4A8 MoE kernels (for example 0x813000fa000), which also appear in a package with a different layout, so they are not addresses. A check like this at export time would have rejected the package. We can share the script if it is useful.

Setup

  • 4x DGX Spark GB10 (SM121), driver 615.71.09, CUDA 13.3, PyTorch 2.14.1+cu132 (from the pin's lock file)
  • Hyperflare 7f534bb6, DeepSeek-V4.1-Flash-1017 official checkpoint, TP=4, --roce-ring
  • Export: --prefill-chunk-size 8192 --max-concurrency 6 --mtp-draft-tokens 5 --expert-scale-compression csf
Lingua principale
Rust
Stelle
800
Fork
106
Merge medio
1g 3h
PR unite (30g)
100

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di NVIDIA/hyperflare-inference

Tutte le issue di NVIDIA/hyperflare-inference

Issue simili

Altre issue su Rust

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.