[BCG][GLM5] perf: BCG support and prefill enhancements - #27053
Conversation
…ed eager submodules, improving performance. Added a new composite operation `mla_bmm_then_unified_attention`.
…ved documentation and assertions for `num_tokens`.
… ensuring correct handling of adjacent split operations and their interactions with non-split nodes. Introduced helper functions for graph creation and target extraction.
…_ENABLE_PCG_DSA_EAGER_FUSION` environment variable. Updated `dsa_indexer` and `mla_bmm_then_unified_attention` to conditionally register split operations based on this setting, enhancing performance for CUDA environments.
…LE_PCG_DSA_EAGER_FUSION` environment variable to the test configuration
This reverts commit cefa693.
This reverts commit 7576d1f.
This reverts commit 2feaaef.
# Conflicts: # python/sglang/srt/layers/attention/dsa/dsa_indexer.py
- split op reuses compiled _get_logits_head_gate instead of an inlined copy - rename dsa_indexer_pcg_dispatch -> dsa_indexer_graph_dispatch (serves PCG + BCG)
|
Thanks @Oasis-Git. Agreed that the previous version made this look more complex than it needed to be. The dispatch is intentional, but I changed the design so it is no longer hidden behind an env flag. Under PCG/BCG graph capture this split-op path is now the default because it was the best-performing option in benchmarks. In particular, the indexer can make the k-only decision eagerly inside the split op; for small OSL this avoids running the full path and reduces the number of kernels. So the split op is not an extra eager fast path for the normal execution path. It replaces the in-graph non-CP prefill path only:
I also cleaned up the integration according to your readability comments: inlined the context helper, moved the PCG/BCG split-op functions to the bottom of the file, and named the BCG wrappers explicitly. |
Oasis-Git
left a comment
Oasis-Git
left a comment
There was a problem hiding this comment.
In general it is good to me now. I suggest to move the bcg part in another pr so we can think about how to fix this gap more gracefully. Thanks for your contribution
There was a problem hiding this comment.
It is clean to me. This file should be done.
| self.cuda_graph._segments.append(graph) | ||
|
|
||
| def _end_current_segment(self) -> None: | ||
| def _end_current_segment(self, drop_empty: bool = False) -> bool: |
There was a problem hiding this comment.
From my perspective I think in this pr we do not need to implement this function or at leaset, split it into next pr. I am still thinking about the implementation code style.
|
cc @Fridge003 for mla file check. |
|
/rerun-failed-ci |
4 similar comments
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
/rerun-failed-ci |
|
Hi @mattteochen can this work on GLM5.2 model? |
|
/rerun-failed-ci |
@Fridge003 yes, it works with GLM5.2 by setting Perf gain is replicated: 1 batch prefill bench 1k1k |
Reconcile the DSA indexer dispatch with main's #27053 (BCG + graph split-op-surface prefill): make pcg/bcg_dsa_indexer_prefill_split and the _fused_q/_fused_k store helpers fusion-aware (num_tokens/out_cache_loc graph contract, single-stream under capture), adopt main's logits_head_gate_graph / in_piecewise_or_breakable_cuda_graph, and drop the now-unused k_cache_and_topk_result split op. Fusion-off behavior is identical to main.
|
cc @nvjullin |
4 participants
Motivation
#23351 introduced PCG for GLM5 (
nsa_indexer). As the current implementation runs the full indexer path, which is not fully captured by a CUDA graph, we found that running the fast indexer in eager mode yields better performance.Moreover, the current PCG split creates a CUDA graph island with only 1 kernel, introducing unnecessary overhead during dispatch.
After
2048tokens, PCG is disabled; the full path is being run eagerly in any case.Modifications
bmm+unified_attention_with_outputto avoid single cuda graph with onlybmm(opt in with env variable)Accuracy Tests
PCG:
BCG:
Speed Tests and Profiling
Tested on B200 and INTEL(R) XEON(R) PLATINUM 8570.
1k1k
Pareto for 1k1k on B200 TP8: ~4% speedup over the PCG solution at max concurrency:
Median TTFT for 1k1k sees 10% improvement at max conc (not in pareto).
Benchmark table (pcg is
main, bcg is current branch with the proposed optimizations):8k1k
We also see some small improvement (1% at conc 256) in the high throughput region when extending the maximum number of tokens captured
--piecewise-cuda-graph-max-tokens 8192:Although 8k1k Pareto is not host-bound anymore, we can enable PCG/BCG to achieve stream overlap for MoE:
This distribution, incorporating the full sweep of concurrency
[4, 8, 16, 32, 64, 128, 256]shows how capturing max tokens8192is enough to trigger the BCG path for most of the prefill batches (memory overhead between default value 2k and 8k is 2.71 GB/GPU):BCG empty cuda grah removal
Removing empty cuda graph islands brings very little perf gain, at least for GLM5 where there's only 1 of such case:
A single batch 1k1k benchmark:
Prefill. latency: 0.05460 --> Prefill. latency: 0.0540It's useful to hide empty graph warnings:
/usr/local/lib/python3.12/dist-packages/torch/cuda/graphs.py:126: UserWarning: The CUDA Graph is empty. This usually means that the graph was attempted to be captured on wrong device or stream. (Triggered internally at /pytorch/aten/src/ATen/cuda/CUDAGraph.cpp:147.) super().capture_end()Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ✅ Run #28023999062
Latest PR Test (Extra): ✅ Run #28126718160