[REGRESSION][dev] DeepSeek-V4 CSA repeated recompilation on forward for THD packed sequences.
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the regression**
@NVIDIA/mcore-oncall After #7064, the performance for Deepseek-v4 with THD packing is ~10x slower due to repeated recompilation on the forward hot path on B200.
**To Reproduce**
~With THD packing, `total_k_rows` will change every for every samples and will trigger a host-side Dynamo recompiles.~
Edit: this is not related, please see the linked PR for actual fix.
```python
import time, torch
from megatron.core.transformer.experimental_attention_variant.csa_utils.fused_sparse_attention import (
build_thd_compact_k_layout, pack_thd_compact_k,
)
cu_q = torch.tensor([0, 512, 1024, 2048], dtype=torch.int32, device="cuda")
cu_k = torch.tensor([0, 128, 256, 512], dtype=torch.int32, device="cuda")
torch.cuda.synchronize()
t0 = time.perf_counter()
for total_k in range(500, 540):
cu, m = build_thd_compact_k_layout(cu_q, cu_k, total_k, 4)
pack_thd_compact_k(torch.randn(total_k, 64, device="cuda", dtype=torch.bfloat16), m)
torch.cuda.synchronize()
print(f"{(time.perf_counter()-t0)*1e3:.0f} ms") # before: ~1200 ms; after remove torch.compile: ~200 ms
```
**Previous performance**
A dummy 2B DS-v4 arch with THD seq len 23040, 8 gradient accumulation step.
**Stack trace/logs**
If applicable, add the stack trace or logs related to the regression.
**Environment (please complete the following information):**
- I'm using megatron-bridge but with the latest MCore's dev branch.
- New Megatron-LM commit ID: #7064 after this refactor.
**Proposed fix**
Removing the torch.compile for `build_thd_compact_k_layout` and `pack_thd_compact_k` partially recover the performance, but not entirely due to other triton kernel recompile for thd pack.
**Additional context**
Agent's note:
@torch.compile on build_thd_compact_k_layout and pack_thd_compact_k was added earlier in 0f9c777f7 (#5992, compact BF16/MXFP8 DSA indexer on SM100+).
Intent: fuse the small torch logic (segment padding, row map, index_select) into one compiled graph when shapes are stable — e.g. CUDA graph capture with fixed buffer sizes.
#7064 amplified the cost by making the fused training forward call the compact indexer path inside every layer’s FusedCSAIndexerLossFunc.forward whenever thd_compact_indexer_available (B200 + bf16).
k_indexer_flat.shape[0] is total_k_rows. With packed sequences, that integer changes microbatch to microbatch. Dynamo treats it as a compile-time specialization → new graph per distinct value → host-side compile storm × num_layers per step.
Before #7064, the same compiled functions existed but ran less often on the hot path (e.g. compact layout mainly tied to CUDA-graph workspace prep in Path B, while index combining still used window_idxs + torch). After #7064, Path B + compact Top-K + sequence-major indices became the default fused training pipeline, so the compile decorators landed on the hottest call site.
Contributor guide
Research direction
Start with megatron/core/transformer/experimental_attention_variant/csa_utils/fused_sparse_attention.py, especially build_thd_compact_k_layout and pack_thd_compact_k, then run the supplied THD benchmark. Verify whether changing total_k causes repeated compilation in the forward path and whether the completed fix removes the compile storm while restoring the reported performance.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100