HSTU Thread-Map Presets — Gantt + Stream/Thread/Lookahead 布局

每个 preset 两张图：上图 Gantt（一次 progress() 在各线程上的时间分布，大致 CPU 时间 + 跨线程同 la 等待）；下图拓扑（task → stream / thread / lookahead 静态分配）。
6-la cascade（pipeline.py:386-391）：h2d@5 → start_shuffle@4 → finish_shuffle@3 | start_input_dist@3 → wait_input_dist@2 | prefetch@1 | fwd/bwd/opt@0。

Stream 颜色 default memcpy data_dist prefetch | 线程颜色 io compute data_dist prefetch tNN（per-task）
lookahead: la=k 表示该 task 比同 progress 的 la=0 task 提前处理 batch index +k；la=0 = 处理本步 batch。
Gantt 中的红色虚线：跨线程同 progress 数据/控制依赖中**导致 wait** 的 binding 边。规则：ord 类（slot edge / depends_on）只在 producer.la == consumer.la 时画 wait；不同 la 由 engine ring rotation 跨 progress 满足，本 progress 内**不**画 wait。sync 类（same_progress_sync，比如 prefetch=True 时 prefetch_embeddings→backward）不分 la 永远画。
→ 例：finish_shuffle → start_input_dist 在 d3 (2 vs 1) 下不画，**只有 d6 cascade (3 vs 3 同 la) 才画**。所以前 7 个 d3 variant 没有 io→compute 的同 progress wait，只有 default_d6 / full_split_d6 有。
la=0 cluster 内部（zero_grad / global_tokens / compute_output_dist / forward → backward → finalize → opt → watchdog）所有任务都同 la=0，跨线程时画 wait（per_task 下 forward→backward 把 backward 锁在 forward 之后）。
Gantt task 宽度：粗略 CPU host-thread 时间估计（HSTU H100 单节点，prefetch=True，仅做相对比较；GPU 实际时间另算）。

variant	la cascade	thread_map
legacy_none	d3	— (legacy)
default_d3	d3	io / compute (2)
by_stream_d3	d3	by_stream (4)
per_task_d3	d3	per_task (14)
io_prefetch_compute_d3	d3	io / prefetch / compute (3)
io_data_dist_compute_d3	d3	io / data_dist / compute (3)
full_split_d3	d3	io / data_dist / prefetch / compute (4)
default_d6	d6	io / compute (2)
full_split_d6	d6	io / data_dist / prefetch / compute (4)

📡 NCCL 通信清单 — 实测每 task 内逐个 NCCL 内核（nsys 采样）
下表来自实测 nsys 采样（rank0, 50 iter, HSTU benchmark balanced_ranking_benchmark.gin: local_batch=32, item_emb_dim=128, num_layers=10, hidden=1024, max_seqlen 4096 zipf α=1.2; 8×H100, DP=8, no TP/PP）。一个 task 内可能包含多个 NCCL 内核——本表把每个 NCCL kernel 单独列出（CUPTI_ACTIVITY_KIND_KERNEL shortName 直接采集），并标注它们运行在哪条 GPU stream（streamId 经 cross-check 得到 str-7=default / str-23=TorchRec output_dist comm / str-24=batch_shuffler memcpy）。
Process group：HSTU benchmark 是纯 DP-8，DistributedModelParallel(env=ShardingEnv.from_process_group(WORLD)) + Megatron DP 也是 WORLD（commons/distributed/sharding.py:282-283）→ 所有 NCCL 共用同一个 comm， _NcclOrderedLock ticket 链跨表中所有 NCCL kernel 全局生效。

#	task	stream	CUDA streamId	NCCL kernel (shortName)	payload — 实际是什么 / 量级
1	`start_shuffle`	memcpy	24	`ncclDevKernel_AllGather_RING_LL`	workloads = int64[LBS=32] → 全局 256, ~2 KB，1 channel = latency-bound
2	`finish_shuffle` → `pad_and_allgather_batch` 内 KJT lengths AllGather + KJT values AllGatherV，2 个 kernel 串行	memcpy	24	`ncclDevKernel_AllGather_RING_LL`	KJT lengths：fixed-size int64×(LBS×#feature) → 全局 [256×#feat]，~10 KB，1 channel
2		memcpy	24	`ncclDevKernel_AllGatherV_RING_LL`	KJT values：variable-size int64 sum-of-seqlen，~MB，24 channel = bandwidth-bound
3	`start_input_dist` → TorchRec splits a2a，2 pipelined module 融合成 1 个 grouped call	data_dist	23	`ncclDevKernel_SendRecv`	splits header (lengths-of-lengths) — int32×8×#feat，~hundreds B；8 channel = 8-rank a2a fanout
4	`wait_input_dist` → `splits_awaitable.wait()` 在 sharded `EmbeddingCollection`（仅 1 个 model-parallel 模块）上触发底层 `KJTAllToAllTensorsAwaitable`，它把 KJT 的 lengths 和 values 拆成 2 个独立的 a2a NCCL group，串行发射	data_dist	23	`ncclDevKernel_SendRecv` (KJT lengths a2a)	每条样本在每个 feature 上的 token 数 — int64×#feat×LBS 的拆分，~tens of KB，按 8-rank a2a fanout
4		data_dist	23	`ncclDevKernel_SendRecv` (KJT values a2a)	实际的 token id（int64 sum-of-seqlen × #feat），按 feature shard 重分布到 owner rank，~MB
5	`global_tokens_allreduce`	default	7	`ncclDevKernel_AllReduce_Sum_f32_RING_LL`	scalar int64 token count → 1 channel，~8 B（kernel cast 成 f32 SUM），latency-bound
6	`compute_output_dist` → embedding output a2a (forward)	default ⚠ 但 NCCL kernel 落到 str-23（TorchRec 用自己的 output_dist comm stream 跑 a2a）	23	`ncclDevKernel_SendRecv` (grouped, 2 modules)	embedding output: `(global_batch×Σseq) × 128 × 4B` (FP32 qcomm)，~tens of MB，32 channel = 跑满最大并行
7	`backward` → `.backward()` 内 autograd 触发仅有 embedding output dist 的 a2a 反向；不含任何 AllReduce — gradient AllReduce 全部在 `finalize_model_grads` 里发射	default ⚠ NCCL kernel 落到 str-23（与 #6 同 comm stream），并且由 `pt_autograd_0` 线程 launch（不是 engine 的 python 主线程），所以不被 `[engine] backward` NVTX 包裹	23	`ncclDevKernel_SendRecv` (emb output dist a2a backward, launched from pt_autograd_0)	embedding output 的反向 a2a → ~tens of MB，p99 ~26 ms 是 critical-path gate 后被前序 NCCL 阻塞所致（疑似根因）
8	`finalize_model_grads` → 所有 gradient AllReduce 都在这里发射（HSTU benchmark `overlap_grad_reduce=False`，backward 不嵌入 grad bucket AR）；本 benchmark 配置下实测 2 个 AllReduce/iter（Megatron 把 dense param grads 切成 2 个 bucket 串行发射）	default	7	`ncclDevKernel_AllReduce_Sum_f32_RING_LL` (gradient bucket 1)	HSTU dense param grads bucket 1，~tens of MB，9 channel
8		default	7	`ncclDevKernel_AllReduce_Sum_f32_RING_LL` (gradient bucket 2)	HSTU dense param grads bucket 2，~MB，9 channel

采集方法：rank0 sqlite (nsys export) 上跑 SQL 把每个 [engine] <task> NVTX 范围与CUPTI_ACTIVITY_KIND_RUNTIME → CUPTI_ACTIVITY_KIND_KERNEL（按 correlationId 配对）做时间包含连接，再按 shortName LIKE 'ncclDevKernel%' 分组得到每 task 内的 NCCL kernel 列表 + GPU duration + gridX (= NCCL channels)。 backward 那一行的 SendRecv 因为是从 pt_autograd_0 线程 launch 的，[engine] backward 包不到 — 我们改用按 thread name 过滤来归因。
关键观察： (a) 每 iter 共 11 个 NCCL kernel（1 AG + 1 AG + 1 AGV + 1 + 2 SendRecv + 1 AR + 1 SendRecv + 1 SendRecv + 2 AR），跨 3 条不同的 GPU stream（7 / 23 / 24）； (b) compute_output_dist 和 backward 的 emb a2a 反向都跑在 str-23（TorchRec 的 output_dist comm stream），不是 default； (c) HSTU benchmark 没开 overlap_grad_reduce，所以 backward 内只剩 emb a2a 反向 1 个 SendRecv，dense param AllReduce 全部 defer 到 finalize_model_grads； (d) start_input_dist kernel duration 长 = NCCL ticket 等待前序 collective（_NcclOrderedLock 串行链 + 跨 rank 同步等待），不是数据量本身大。

🧵 代码 stream ↔ nsys streamId 映射 + 任务跨 stream 拓扑

代码层 stream	nsys streamId
`default`	7
memcpy slot	24
data_dist slot	28
prefetch slot	32
TorchRec output_dist comm	23
`ShardedEmbedding._side_stream`	27

Task	stream(s)
`start_input_dist`	data_dist (28) + comm (23) NCCL
`wait_input_dist`	data_dist (28) + comm (23) NCCL
`prefetch_embeddings`	data_dist (28) → prefetch (32)
`compute_output_dist`	default (7) + comm (23) NCCL
`backward` (emb a2a)	default (7) + comm (23) NCCL
`ranking_embedding_forward`	default (7) + _side_stream (27) if OVERLAP_DP_MP=1
`start_shuffle` / `finish_shuffle`	memcpy (24)
`global_tokens_allreduce` / `finalize_model_grads` / `optimizer_step` / `forward`	default (7)

源数据： pipeline.py（HSTU_DEFAULT_THREAD_MAP / HSTU_THREAD_MAP_PRESETS / 6-la cascade lookaheads） · tasks.py（每个 task 的 stream + reads/writes）。 Gantt 用的 task CPU 估时是粗略数（HSTU H100 单节点，prefetch=True），相对长度比绝对值更有意义；GPU 真实执行时间另算（fwd/bwd 在 default stream 上 ~25ms 各自占主导）。

HSTU Thread-Map Presets — Gantt 时间轴 + Stream / Thread / Lookahead 拓扑