[FEA] Refine training benchmark model configuration.

[JacoCheung]

Currently, our benchmark configuration can not reflect real scenarios well.

1. The batchsize or sequence length is not typical
2. Input data is uniform, we need powerlaw. (both seqlen and key range)
3. The embedding dim is equal to hidden dim
4. The attention hyper-params is not well set
5. The model weight is not large enough. (because of 3. and 4.)
6. We have only single node test.
7. We do not benchmark the TP.
   Therefore we need to improve our benchmark suites.

---

By submitting this issue, you agree to follow our code of conduct and our contributing guidelines.

[JacoCheung]

Hi @shijieliu . I would like issue #117 to address hstu layer perf only. We can move to this issue regarding model perf benchmark.

usually should be 1k - 8k

I think batchsize might be coupled with sequence length. If we use larger hidden_size and batchsize X seqlen. It may lead to OOM easily. In that sense, I think we need to add large scaling (like 16x8 GPUs) benchmark (or with real datasets).

I think we can use larger hidden_size
You mean we can use more neads or larger head_dim?

Does this benchmark reflect embedding as well? (i.e. use small dense but intensive dyanmic embedding).

[JacoCheung]

Update with more specfic requirements.

1. Setup a baseline: Sufficient embeddings + dense (at least 2 nodes)
   1. Embedding all offloaded to CPU.
   2. (need to balance the workload ration between sparse vs dense)
   3. Random dataset on CPU.
2. Turn on op level opt.
   1. HSTU cutlass
   2. Embedding Cache
   3. recompute
   4. TP
3. Pipeline (Prefetch)
4. Workload balancer

[shijieliu]

junjie is working on figuring out the suitable workload(seqlen distribution) and after verifying the optimization effect, will dig into how to hide workload balancer with prefetch.

[shijieliu]

validate the effectiveness of latest dynamicemb prefetch

[JacoCheung]

Regarding the issue that shuffle plus pipelining fails to overlap communication and computation, I did an in‑depth analysis and found that the core problem is a large number of syncs on the main thread. These syncs include:

1. D2H/H2D copies, involving all2all / allgatherv and the KK algorithm (KK is pure CPU). Because the output buffer needs real‑time size information, some sizes can be computed ahead of time (for example, the case just fixed in triton split ), but some sizes have to be computed on the critical path. In addition, pageable memory may also lead to longer sync times.

2. JIT compilation, including Triton / nvJet. This is very hard to avoid under dynamic shapes.

All of these syncs block the CPU thread and prevent NCCL kernels from being launched quickly, so communication cannot fully overlap with computation. As a result, the final performance is actually worse than running without pipelining.