NVIDIA / NVIDIA/Megatron-LM

[REGRESSION][dev] DeepSeek-V4 CSA repeated recompilation on forward for THD packed sequences.

Open
#7,340 0 comments 0 reactions 0 assignees View on GitHub
community-request
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.

Image

**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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.