Motivation

This PR fixes two bugs when running on NVIDIA B200 (Blackwell) GPUs with piecewise CUDA graph enabled:

bug 1: Missing forward metadata initialization

Problem: AttributeError: 'DecodeMetadata' object has no attribute 'prefill_wrappers' occurs during warmup phase in piecewise CUDA graph runner.

Root Cause: The warmup_and_capture() method in piecewise_cuda_graph_runner.py calls model.forward() directly without initializing the attention backend's forward metadata, causing the metadata to be uninitialized or in an incorrect state.

Fix: Add self.model_runner.attn_backend.init_forward_metadata(forward_batch) before calling model.forward() in the warmup phase.

bug 2: CUTLASS backend TMA descriptor initialization failures on B200

Problem: When using CUTLASS backend with piecewise CUDA graph on B200, the server crashes with Error: Failed to initialize the TMA descriptor errors showing invalid tensor dimensions (e.g., globalDim (128,0,8,1,4) with zero dimension).

Root Cause: FlashInfer's CUTLASS backend has compatibility issues with B200's TMA (Tensor Memory Accelerator) features when combined with piecewise CUDA graph capture.

Fix: Intelligently disable CUTLASS backend and fall back to FlashAttention-2 only when both B200 (SM100) and piecewise CUDA graph are enabled. A clear warning message informs users of this backend selection. When piecewise CUDA graph is disabled, B200 can still use CUTLASS for optimal performance.

python3 -m sglang.launch_server --model-path Qwen/Qwen3-Coder-30B-A3B-Instruct --tp 4 --host 0.0.0.0 --enable-piecewise-cuda-graph --piecewise-cuda-graph-compiler eager

result:

python3 benchmark/gsm8k/bench_sglang.py --num-questions 2000 --parallel 2000 --num-shots 8
/usr/local/lib/python3.12/dist-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
100%|████████████████████████████████████████████████████████████| 1319/1319 [00:22<00:00, 59.25it/s]
Accuracy: 0.929
Invalid: 0.001
Latency: 22.355 s
Output throughput: 7932.862 token/s

Without fix for Issue 2

100%|████████████████████████████████████████████████████████████| 1319/1319 [00:55<00:00, 23.56it/s]
Accuracy: 0.174
Invalid: 0.817
Latency: 56.065 s
Output throughput: 10410.864 token/s

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.