[Bug] DeepSeek-V4 TBO decode + CUDAGraph: hipBLASLt heuristic during capture aborts with HIP 900 ('operation not permitted when stream is capturing')

[yhl-amd]

Summary

Running DeepSeek-V4 with decode-side TBO (--enable-tbo all) under the default level-3 (piecewise + CUDAGraph) compilation aborts during CUDAGraph capture with:

Hip error: 'operation not permitted when stream is capturing'(900) at hipblaslt.cpp:147

repeated across ranks, followed by proc died / Received unexpected SHUTDOWN signal from DP rank N during initialization. The server never reaches ready.

This is a TBO + hipBLASLt + CUDAGraph interaction, independent of the attention/KV implementation. --enforce-eager + decode-TBO works fine; only the CUDAGraph capture path is affected.

Environment

• Model: DeepSeek-V4-Flash-FP8, -tp 4 --data-parallel-size 2 --enable-dp-attention --enable-expert-parallel
• --enable-tbo all (i.e. enable_tbo=True, enable_tbo_decode=True)
• Default compilation: level-3, use_cudagraph=True
• torch 2.10.0+rocm7.2.4, HIP 7.2.53211, gfx950

Isolation

config	CUDAGraph capture
CG + no TBO	✅ ok
CG + --enable-tbo prefill (prefill-ubatch capture only)	✅ ok
CG + --enable-tbo all (decode-ubatch capture)	❌ HIP 900 at hipblaslt.cpp:147

So the failure is specific to decode-ubatch CUDAGraph capture (build_for_cudagraph_capture → _prepare_ubatch_decode at bs > 2 → UBatchWrapper.capture_tbo_graph). The log shows the ubatch half-batch GEMM shape (M:256 for a split of graph bs=512) first appearing at the same moment as the HIP 900 spam.

Root cause

atom/utils/tbo/ubatch_wrapper.py::capture_tbo_graph spins up N ubatch threads, each with its own compute/comm stream, and pre-initializes cuBLAS on both streams before capture:

# atom/utils/tbo/ubatch_wrapper.py ~L297
# Initialize cuBLAS on both streams BEFORE barrier.
with torch.cuda.stream(tbo_ctxs[idx].compute_stream):
    _ = torch.cuda.current_blas_handle()
with torch.cuda.stream(tbo_ctxs[idx].comm_stream):
    _ = torch.cuda.current_blas_handle()

But the crash comes from hipBLASLt (hipblaslt.cpp:147), a different library from cuBLAS. hipBLASLt performs a per-(GEMM shape, stream) heuristic / workspace query (which syncs / allocates) the first time a shape is seen on a stream. The ubatch half-batch GEMM shapes are first encountered on the ubatch threads' streams inside the capture region, so the heuristic query runs during capture → operation not permitted when stream is capturing (HIP 900).

torch.cuda.current_blas_handle() initializes the cuBLAS handle but does not warm hipBLASLt's per-shape algo selection, so the existing pre-warm doesn't cover this path.

Suggested fix

Before capture, run the ubatch model.forward eagerly on the same ubatch compute/comm streams (one dry pass per ubatch shape) so hipBLASLt's per-shape heuristic cache is populated outside capture — analogous to the existing single-stream pre-capture warmup in model_runner but routed through the ubatch streams. Alternatively force a fixed/default hipBLASLt algo (skip heuristic) during capture, or call a hipBLASLt warm-up API for the ubatch shapes.

Note on reproducibility on main

On current main, decode-TBO + CUDAGraph capture crashes earlier — at an unrelated dtype assert in the ubatch decode metadata build:

AssertionError: V4 buffer 'ub0_v4_batch_id_per_token' dtype mismatch: buffer is int64, but got arr with dtype int32.

(_alloc_v4_metadata_buffers allocates the ubatch v4_batch_id_per_token buffer as int64 while the non-ubatch buffer and the staged data are int32.) That assert masks this hipBLASLt issue. Fixing the dtype (int64→int32) lets capture proceed far enough to hit the hipBLASLt-during-capture failure described above. Both were found while validating DeepSeek-V4 paged-SWA under TBO decode.

[yhl-amd]

remove --data-parallel-size --enable-expert-parallel