Tracing the CUBLAS Fault: Debugging EP8 and NVLS on NVLink Blackwell

The Message

[assistant] ## Agent Reasoning
I need to narrow down whether EP8 or NVLS is triggering the CUBLAS failure on the BF16 strided batched gemm operation — this is likely happening in either the MLA attention mechanism or the draft model, and while EP8 changes the MoE routing, the attention computation should still be using TP8. Let me trace through the full crash log to pinpoint exactly where the gemm is failing.
Still crashing. Let me get the full traceback to isolate whether it's EP or NVLS:
[bash] timeout 30 ssh -o ConnectTimeout=15 root@86.38.182.109 "journalctl -u sglang-k26-ddtree-ep.service --no-pager -n 60 | grep -E 'File \"|line [0-9]|Error|cublas|gemm|cuda graph|Capture|draft|verify' | tail -25" 2>&1
May 30 15:02:13 mild-hope-wilts-fin-03 python[33085]:   File "/root/venv_sglang211/lib/python3.12/site-packages/sglang/srt/model_executor/model_runner.py", line 2704, in _dummy_run
May 30 15:02:13 mild-hope-wilts-fin-03 python[33085]:   File "/root/venv_sglang211/lib/python3.12/site-packages/sglang/srt/model_executor/model_runner.py", line 2693, in run_once
May 30 15:02:13 mild-hope-wilts-fin-03 python[33085]:   File "/root/venv_sglang211/lib/python3.12/site-packages/sglang/srt/models/kimi_k25.p...

Context: The Optimization Frontier

This message sits at a critical inflection point in a high-stakes deployment. The assistant has been deploying Kimi K2.6, a Mixture-of-Experts (MoE) language model, with DFlash speculative decoding across two hardware platforms: a PCIe-connected 8× RTX PRO 6000 system (the "PRO6000 box") and an 8× B300 SXM6 NVLink machine. The B300 machine represents the ideal hardware target — 275 GB per GPU, NVLink interconnects with NVLS (NVIDIA's multicast collectives), and sm_103 compute capability. After successfully deploying the DDTree speculative decoding stack on B300 and measuring a 2.15× speedup over the autoregressive baseline at single-request concurrency (285 tok/s vs 133 tok/s), the assistant hit a puzzling observation: the GPUs were showing 100% utilization but drawing only 360W out of a 1100W power budget.

This power discrepancy is the key puzzle. When GPUs are fully utilized but drawing only one-third of their thermal design power, the bottleneck is almost certainly memory bandwidth rather than compute. For an INT4-quantized MoE model like K2.6, this makes sense: each decode step must load expert weights from HBM into the compute units, and the Marlin dequantization kernel is memory-bound. The user astutely suggested trying Expert Parallelism (EP) with NVLink multicast (NVLS) — the hypothesis being that EP would keep more expert weights resident on each GPU, reducing redundant HBM traffic, while NVLS would accelerate the all-to-all communication needed to route tokens between GPUs.

The assistant's response to this suggestion (message 11788) was to create an EP8 + NVLS service variant, adding --ep-size 8 --enable-nccl-nvls --enable-flashinfer-allreduce-fusion to the DDTree configuration. That crashed with a CUBLAS error. The assistant then removed the allreduce-fusion flag (message 11789), reasoning it was the likely culprit, but the EP8 + NVLS combination still crashed. This brings us to the target message (11790), where the assistant must isolate the root cause of the persistent CUBLAS failure.

The Reasoning: A Systematic Debugging Approach

The assistant's reasoning reveals a methodical diagnostic process. The key insight is that the CUBLAS failure occurs on a "BF16 strided batched gemm operation" — a batched matrix multiplication using bfloat16 precision. The assistant correctly identifies two possible sites for this operation: the MLA (Multi-head Latent Attention) mechanism, which is the core attention computation in K2.6, or the draft model's verification path. The reasoning also notes that EP8 changes the MoE routing but should not affect the attention computation, which remains under tensor parallelism (TP8).

This distinction is crucial. In SGLang's implementation, expert parallelism distributes MoE layers across GPUs, with tokens being routed via all-to-all communication. The attention layers, however, remain under tensor parallelism. So if the CUBLAS error is in the attention path, EP8 is not the direct cause — it could be NVLS or some interaction between the two. If the error is in the MoE path, EP8's routing changes could be exposing a bug in the NVLS all-to-all implementation.

The assistant's decision to extract the full traceback from the journal is the correct next step. Rather than guessing or randomly flipping flags, they go to the source: the Python stack trace showing exactly which function call triggered the CUDA error. The grep pattern is carefully chosen — File "|line [0-9]|Error|cublas|gemm|cuda graph|Capture|draft|verify — to capture file paths, line numbers, error messages, and keywords related to the speculative decoding pipeline. The tail -25 limits the output to the most relevant lines.

Input Knowledge Required

To understand this message, the reader needs substantial context about the broader deployment. First, one must understand the hardware topology: the B300 SXM6 machine uses NVLink interconnects that provide significantly higher bandwidth than PCIe, and NVLS (NVIDIA's multicast collective communication library) leverages NVLink's ability to broadcast data to multiple GPUs simultaneously. Second, one needs familiarity with SGLang's parallelism options: tensor parallelism (TP) splits each layer's weights across GPUs, requiring AllReduce after every layer; expert parallelism (EP) keeps complete experts on each GPU and routes tokens via all-to-all communication. Third, the concept of CUDA graphs is important — these are pre-compiled sequences of CUDA operations that eliminate kernel launch overhead, and the assistant has previously confirmed they work for decode steps but not prefills.

The reader also needs to understand the speculative decoding architecture: DFlash uses a small draft model to predict multiple tokens per step, which are then verified by the target model in parallel. The DDTree variant constructs a tree of draft tokens to explore multiple continuation paths. The "verify path" refers to the target model's forward pass over the drafted tokens, which uses a strided batched GEMM — exactly the operation that's failing.

Assumptions and Potential Pitfalls

The assistant makes several assumptions in this message. The primary assumption is that the CUBLAS failure is caused by either EP8 or NVLS specifically, rather than some other configuration change. This is reasonable given that the previous DDTree configuration (TP8 without EP/NVLS) worked correctly, and the only changes were adding --ep-size 8, --enable-nccl-nvls, and (in the first attempt) --enable-flashinfer-allreduce-fusion. However, there's a subtle assumption that the crash is deterministic and reproducible — that the same error will appear on every restart, which the journal output confirms.

Another assumption is that the BF16 strided batched GEMM is occurring in either the MLA attention or the draft model. While these are the most likely candidates, there could be other BF16 GEMM operations in the model — for example, in the gate projections for the MoE layers, or in the output projection after attention. The assistant's reasoning acknowledges this uncertainty by planning to trace the full stack trace.

There's also an implicit assumption that the NVLS implementation in SGLang is compatible with the DDTree speculative decoding path. This is not guaranteed — NVLS may have been tested only with standard autoregressive decoding, and the tree verification step may use GEMM configurations (e.g., specific batch sizes or matrix dimensions) that NVLS doesn't handle correctly.

Output Knowledge Created

This message creates diagnostic output: the partial stack trace showing the error originates in model_runner.py at line 2704 (_dummy_run) and line 2693 (run_once), then enters the model-specific code in kimi_k25.py. The _dummy_run function is particularly telling — it's used for CUDA graph capture, where the engine runs a dummy forward pass to record the sequence of operations. This means the CUBLAS error occurs during graph capture, not during actual inference. This is a critical finding: the CUDA graph capture mechanism is failing when EP8 + NVLS is enabled, which suggests the graph capture doesn't properly handle the EP/NVLS communication primitives.

The truncated traceback at kimi_k25.py doesn't reveal the exact line, but the fact that it reaches the model-specific code narrows the search considerably. The assistant now knows the error is in the model forward pass during graph capture, not in the draft model or verification logic. This shifts the debugging strategy: instead of looking at the speculative decoding code, the assistant should investigate how EP8 + NVLS interacts with the model's forward pass during CUDA graph capture.

The Broader Significance

This message exemplifies the kind of systematic debugging that characterizes high-performance ML engineering. The assistant doesn't panic at the CUBLAS error or randomly try different flag combinations. Instead, they formulate a hypothesis (the error is in MLA or the draft model), identify the information needed to test it (the full stack trace), and execute a precise diagnostic command. The reasoning shows an understanding of the system's architecture — knowing which operations use BF16 strided batched GEMM, understanding the separation between attention and MoE paths, and recognizing that CUDA graph capture is a distinct phase from inference.

The outcome of this diagnostic step will determine the next phase of optimization. If the error is in NVLS's handling of the attention GEMM during graph capture, the assistant might need to disable CUDA graphs for the EP configuration, accepting lower throughput but gaining the benefits of expert parallelism. If the error is in EP8's interaction with the draft model, the fix might involve changes to the speculative decoding pipeline. Either way, this message represents the pivot point where the assistant moves from "try the obvious optimization" to "understand why it broke and find the right fix."