Capturable grammar hangs decode CUDA-graph capture when capture batch x draft tokens exceeds 128 logits rows (xgrammar, TP=8, B300)

[abatilo]

Summary

Enabling --grammar-backend xgrammar (default capturable grammar) makes the server hang forever during decode CUDA-graph capture when the largest capture batch size times --speculative-num-draft-tokens exceeds 128 logits rows. The hang is 100% reproducible (5/5 starts) on an 8x B300 node (TP=8, NVLink with fabric manager / NVLS) and does not reproduce on H100 or H200 nodes whose configurations stay at or under 128 rows. The same B300 configuration is healthy without the grammar backend, and healthy with --grammar-backend xgrammar --disable-capturable-grammar.

Environment

• TokenSpeed commit 76b6295 (head of PR feat(glm-5.3-flash): add day-0 serving support #1259, GLM-5.3-Flash day-0 support)
• Model: zai-org/GLM-5.3-Flash, TP=8, MTP speculative decoding
• 8x B300 (sm_103) on one node, NVLink fabric with fabric manager, NVLS enabled
• torch 2.13.0+cu130, flashinfer 0.6.18 (nightly 20260811 wheels + cubin + jit-cache), xgrammar 0.2.3, CUDA 13.0, driver 580.159.04, Python 3.12

Reproduction

tokenspeed serve \
  --model zai-org/GLM-5.3-Flash --trust-remote-code \
  --tensor-parallel-size 8 --enable-expert-parallel \
  --gpu-memory-utilization 0.85 --max-num-seqs 64 \
  --speculative-algorithm MTP \
  --speculative-draft-model-path zai-org/GLM-5.3-Flash \
  --speculative-num-steps 3 --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 4 \
  --grammar-backend xgrammar

The variable that matters is max-num-seqs 64: capture sizes are derived from it, so the graph list becomes [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 63, 64] and the MTP verify captures are 63 x 4 = 252 and 64 x 4 = 256 logits rows.

What happens

Capture runs largest-first. The bs=64 graph captures fine (~14 s). The bs=63 capture never completes: no further engine log lines, all 8 ranks spin at ~1 core each, and the process never becomes ready (our startup probe eventually kills it; every restart hangs at the same point). The last log lines are always (the FutureWarning is emitted once per rank during the bs=64 warmup forwards):

[AutoTuner]: No tuned config covers trtllm_batch_decode_mla input_shapes=((252, 1, 8, 512), (252, 1, 2052), (252,), (252, 1, 8, 512), (252,)); falling back ...
[AutoTuner]: No tuned config covers trtllm_batch_decode_mla input_shapes=((63, 1, 8, 512), ...); falling back ...
torch/distributed/c10d_logger.py:83: FutureWarning: `torch.distributed.all_gather_into_tensor` is deprecated.

That FutureWarning never appears on the H100/H200 configurations. It comes from LogitsProcessor._get_logits — the full-vocab logits gather uses the Triton RSAG symmetric-memory kernel only up to _LOGITS_AG_MAX_TOKENS = 128 rows (runtime/layers/logits_processor.py) and falls back to torch.distributed.all_gather_into_tensor (NCCL) above that. With a grammar backend enabled that gather is recorded into the captured decode graph. The node's fabric manager logs NVLink multicast team setup/teardown (NVLS registration) in the same minute as every stall.

Working / failing matrix (same image, same model, same flags except where noted):

GPUs	max-num-seqs	rows (x4 draft)	grammar	result
8x H100	16	64	xgrammar	healthy
8x H200	32	128	xgrammar	healthy
8x B300	64	252/256	xgrammar	hangs at bs=63 capture, 5/5
8x B300	64	252/256	none	healthy
8x B300	64	252/256	xgrammar + --disable-capturable-grammar	healthy: all 12 graphs captured in 76 s, NCCL fallback included
8x B300	64 capped: --max-cudagraph-capture-size 32	128	xgrammar	healthy

What we ruled out

• The hostfunc load_inline extension and the xgrammar Triton bitmask kernel: the bs=64 graph, which contains both, captures and completes.
• A Triton compile problem for a non-multiple-of-16 batch: the Hopper configurations capture bs=15/31 with the same kernels.
• The NCCL fallback inside capture on its own: with --disable-capturable-grammar (eager grammar buffers, everything else identical) the same 252/256-row graphs capture and the server is healthy.

So the hang needs both the capturable grammar path (CapturableGrammarExecutor: the side stream forked from the captured forward plus the cudaLaunchHostFunc nodes) and the >128-row graphs whose logits gather is the c10d/NCCL fallback rather than the RSAG symmetric-memory kernel. The bs=64 graph, which already contains both, captures; the second such graph (bs=63) is the one that never returns. We have not captured a stack of the hung ranks, so we cannot say which call blocks.

Workaround

Either --max-cudagraph-capture-size 32 (keeps every captured graph at or under 128 rows, so the gather stays on the RSAG kernel; batches above 32 run eager) or --disable-capturable-grammar (keeps full-size capture, loses the overlapped grammar fill).

Possible fixes

• Size _LOGITS_AG_MAX_TOKENS from the capture configuration so captured graphs never take the c10d fallback when a grammar backend is on, or
• make the capturable grammar path safe alongside the c10d fallback (or refuse that combination with a clear error instead of hanging).