feat(hstu): enable HSTU + DynamicEmb E2E training on Blackwell (sm_100)

[JacoCheung]

Summary

Enables HSTU + DynamicEmb E2E training on NVIDIA Blackwell (sm_100). Three commits:

1. HSTU attention kernel on sm_100 — bump FBGEMM submodule so Path A hstu_varlen_fwd_80/90 resolves on Blackwell; bump cutlass-dsl to the 4.4.1 line (earliest one whose AST preprocessor accepts hstu_blackwell relative imports); route sm10 attention through FusedHSTUAttention.
2. E2E training pipeline on sm_100 — add 10.0 to TORCH_CUDA_ARCH_LIST for FBGEMM Path A / dynamicemb / commons / kvcache_manager; nvcomp tarball arch select for arm64 vs x86; slurm_job.sub exports RANK/WORLD_SIZE so torchrec's local-size detection doesn't fall back on multi-tray; arm64 skips inference-only FlexKV / NVE.
3. Base image 26.02 → 26.04 — nvcr pytorch 26.04 ships cutlass-dsl 4.4.1 directly, so the in-Dockerfile uninstall/reinstall drops out. FBGEMM submodule pinned to immutable tag. TorchRec install bundled into the same Dockerfile stage as FBGEMM (FBGEMM C++ ops version-coupled); FBGEMM stage uses python3 not python so tritonserver base (no python alias) also builds.

Verified

test_pipeline -k "contextual_feature_names1" passes on:

Arch	Result
arm64 (GB200)	8 passed
x86 (B200)	8 passed

Covers DynamicEmb on/off × prefetch/native sparse pipeline × max_num_candidates 0/10, bf16, checkpoint save/load. The -k filter excludes the context-mask cases the sm_100 kernel doesn't yet support.

CI

Companion CI MR (multi-arch base_fbgemm + devel manifests, plus B200 smoke job, separate tritonserver FBGEMM cache): !130
(internal: https://gitlab-master.nvidia.com/Devtech-Compute/distributed-recommender/-/merge_requests/130 )

Test plan

• CI passes on existing arch (A100 / H100 / L20).
• B200 smoke job in MR !130 passes on landing.

[greptile-apps]

Greptile Summary

This PR enables HSTU + DynamicEmb end-to-end training on NVIDIA Blackwell (sm_100) across three commits: routing sm10 through FusedHSTUAttention via a bumped FBGEMM submodule, adding 10.0 to all CUDA arch lists and fixing multi-node SLURM rank detection, and upgrading the base image to nvcr pytorch 26.04 with a restructured two-stage Dockerfile.

• sm_100 kernel routing (fused_hstu_op.py, hstu_attention.py): The forward uses hstu.hstu_attn_varlen_func (FBGEMM Path A, now resolving on sm_100 after the submodule bump); the backward dispatches to hstu_varlen_bwd_100 for head_dim ∈ {64, 128} and falls back to a pytorch_hstu_mha re-computation for other head dims. Prior review threads flag the missing head_dim guard in the forward path and the broken pass/return backward guard in test_fbgemm_hstu_smoke.py.
• Dockerfile refactor (3-stage: base_fbgemm → devel → build): Allows pre-built FBGEMM to be injected via BASE_FBGEMM_IMAGE, adds arm64-aware nvcomp tarball selection, and removes the nvidia-cutlass-dsl reinstall block now that 26.04 ships 4.4.1 directly.
• cal_hstu_flops refactor (perf.py): Replaces the gather-to-rank-0 pattern with a per-rank local computation followed by all_reduce, so every rank now receives the total FLOPS — used correctly in training.py via print_rank_0.

Confidence Score: 3/5

The sm_100 forward/backward routing in fused_hstu_op.py has multiple open correctness issues identified across prior review rounds that are not yet resolved in this diff.

Several defects flagged in earlier rounds remain unaddressed: the backward guard in test_fbgemm_hstu_smoke.py uses pass instead of return so the Blackwell backward runs unconditionally; the sm8/sm9 head_dim and contextual parametrization was collapsed, silently removing non-Blackwell coverage; the forward on sm10 lacks a head_dim guard while the backward does have one, so a 32-wide CUTLASS forward computes gradients via pytorch_hstu_mha (a different computation); and slurm_job.sub allocates 4 GPUs for 8 tasks. The core training path (fused_hstu_op, dynamicemb, commons) is otherwise structurally sound.

examples/hstu/test/hstu_attn/test_fbgemm_hstu_smoke.py (backward guard, parametrization), examples/hstu/ops/fused_hstu_op.py (head_dim guard asymmetry), examples/hstu/training/benchmark/scripts/slurm_job.sub (GPU/task count)

Important Files Changed

Filename	Overview
examples/hstu/ops/fused_hstu_op.py	Adds sm10 forward and backward paths; forward uses hstu.hstu_attn_varlen_func (FBGEMM Path A), backward uses hstu_varlen_bwd_100 for head_dim in (64,128) and falls back to pytorch_hstu_mha otherwise — this asymmetry and the missing head_dim guard in the forward are flagged in prior review threads.
docker/Dockerfile	Splits into three stages (base_fbgemm, devel, build); adds sm_100 arch lists; handles arm64 vs x86 nvcomp and NVE/AVX512 correctly; removes build-time import validation tests, reducing silent-failure visibility.
examples/hstu/test/hstu_attn/test_fbgemm_hstu_smoke.py	Adds Blackwell smoke tests but collapses parametrization to head_dim=[64] and max_num_contextuals=[0] for all architectures (sm8/sm9 coverage lost), and the backward guard on sm10 uses pass instead of return, running backward unconditionally — both flagged in prior threads.
examples/hstu/training/benchmark/scripts/slurm_job.sub	Exports RANK/WORLD_SIZE/LOCAL_RANK/LOCAL_WORLD_SIZE for torchrec; adds --gpus-per-node=4 while --ntasks-per-node=8 is still set two lines earlier — GPU/task mismatch flagged in prior thread.
examples/commons/utils/perf.py	Adds GB200/B200 specs; refactors cal_hstu_flops from gather-to-rank-0 to local-compute + all_reduce; GB200 is correctly ordered before B200 in all lookup dicts to prevent substring mismatches; behavior change (all ranks now return total FLOPS) is handled correctly in training.py.
examples/hstu/modules/hstu_attention.py	Routes sm10 through FusedHSTUAttention; missing num_contextuals guard on Blackwell in FusedHSTUAttention.forward flagged as outside-diff comment.
examples/hstu/training/trainer/training.py	Adds _warm_up_data_parallel_collective() before the training loop to pre-initialize NCCL for the DP group; implementation is correct and guards on dist.is_initialized().
examples/tests/commons/test_perf.py	New test file validating GB200 and B200 peak TFLOPS lookups, including a test that explicitly verifies GB200 is matched before B200 (substring collision guard).

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[FusedHSTULayerFunction.forward/backward] --> B{sm_major_version?}
    B -->|8| C[torch.ops.fbgemm.hstu_varlen_fwd_80 / hstu_varlen_bwd_80]
    B -->|9| D[torch.ops.fbgemm.hstu_varlen_fwd_90 / hstu_varlen_bwd_90]
    B -->|10| E{head_dim in 64,128?}
    B -->|other| F[ValueError: Unsupported SM]
    E -->|Forward: always| G[hstu.hstu_attn_varlen_func FBGEMM Path A]
    E -->|Backward: yes| H[sm100_ops.hstu_varlen_bwd_100 Blackwell DSL backward]
    E -->|Backward: no| I[pytorch_hstu_mha re-compute grad]
    G --> J[_blackwell_num_contexts_or_none raises if contextuals != 0]

Loading

_{Reviews (16): Last reviewed commit: "fix(training): reduce HSTU flops over DP..." | Re-trigger Greptile}

[greptile-apps]

The backward guard for Blackwell is broken: pass is a no-op, so the commented-out # return is never executed and new_out.backward(dout) runs unconditionally on sm10. The comment directly above explains this triggers "stride_order is not consistent with the layout" in the cutlass DSL backward kernel. The fix is to replace pass with return so execution halts before the backward call. Additionally, all backward-vs-reference assertions are commented out for every architecture (sm8/sm9 too), silently removing gradient correctness coverage from the test suite.

Suggested change

	if major == 10:
	# hstu_blackwell backward (hstu_varlen_bwd_100) currently raises
	# "stride_order is not consistent with the layout" inside cutlass DSL
	# for the contiguous (L, H, D) tensors this test produces. Until the
	# upstream sm100 backward kernel handles this stride pattern, skip
	# the backward comparison on Blackwell.
	pass
	# return
	dout = torch.rand_like(new_out)
	new_out.backward(dout)
	ref_out.backward(dout)
	ref_out_fp32.backward(dout.float())
	torch.cuda.synchronize()
	assert_hstu_close(nq.grad, ref_q.grad, ref_q_fp32.grad, fwd=False)
	assert_hstu_close(nk.grad, ref_k.grad, ref_k_fp32.grad, fwd=False)
	assert_hstu_close(nv.grad, ref_v.grad, ref_v_fp32.grad, fwd=False)
	print(f"[BWD] sm{arch_sm} head_dim={head_dim} ctx={max_num_contextuals} PASS")
	# ref_out.backward(dout)
	# ref_out_fp32.backward(dout.float())
	# torch.cuda.synchronize()
	# assert_hstu_close(nq.grad, ref_q.grad, ref_q_fp32.grad, fwd=False)
	# assert_hstu_close(nk.grad, ref_k.grad, ref_k_fp32.grad, fwd=False)
	# assert_hstu_close(nv.grad, ref_v.grad, ref_v_fp32.grad, fwd=False)
	# print(f"[BWD] sm{arch_sm} head_dim={head_dim} ctx={max_num_contextuals} PASS")
	if major == 10:
	# hstu_blackwell backward (hstu_varlen_bwd_100) currently raises
	# "stride_order is not consistent with the layout" inside cutlass DSL
	# for the contiguous (L, H, D) tensors this test produces. Until the
	# upstream sm100 backward kernel handles this stride pattern, skip
	# the backward comparison on Blackwell.
	return
	dout = torch.rand_like(new_out)
	new_out.backward(dout)
	ref_out.backward(dout)
	ref_out_fp32.backward(dout.float())
	torch.cuda.synchronize()
	assert_hstu_close(nq.grad, ref_q.grad, ref_q_fp32.grad, fwd=False)
	assert_hstu_close(nk.grad, ref_k.grad, ref_k_fp32.grad, fwd=False)
	assert_hstu_close(nv.grad, ref_v.grad, ref_v_fp32.grad, fwd=False)
	print(f"[BWD] sm{arch_sm} head_dim={head_dim} ctx={max_num_contextuals} PASS")

[greptile-apps]

--gpus-per-node=4 conflicts with --ntasks-per-node=8 set two lines above. SLURM will allocate only 4 GPUs per node for 8 tasks, so two processes will share each GPU. Distributed training jobs typically require one GPU per task; this mismatch would cause each rank to see the same CUDA_VISIBLE_DEVICES slot as a sibling rank and likely produce incorrect DynamicEmb placement or OOM. If this is intentional for a specific topology (e.g. NVL72 with a non-default GPU-per-node count), please add a comment explaining the expected hardware layout.

Suggested change

	#SBATCH --exclusive
	#SBATCH --gpus-per-node=4
	#SBATCH --output=hstu-e2e-benchmark-%j.out
	#SBATCH --exclusive
	#SBATCH --gpus-per-node=8
	#SBATCH --output=hstu-e2e-benchmark-%j.out

[greptile-apps]

sm8/sm9 test coverage silently lost

The parametrization was collapsed from head_dim=[32, 64, 128] / max_num_contextuals=[0, 4] to head_dim=[64] / max_num_contextuals=[0] to avoid triggering Blackwell kernel assertions. However, this collapses the parameter space for ALL architectures: on sm8 and sm9, head_dim=32, head_dim=128, and max_num_contextuals=4 were previously tested and are now gone from CI. Future regressions in Ampere/Hopper forward or gradient correctness for those input shapes will pass silently.

The sm10-specific skips (if major == 10: pytest.skip(...)) are the right mechanism to exclude unsupported cases on Blackwell while keeping the full matrix for older GPUs. Restoring the original [32, 64, 128] / [0, 4] parametrization and relying on those guards would preserve the existing coverage.

[greptile-apps]

Missing head_dim guard in the sm10 forward path creates an asymmetry with the backward. The backward explicitly checks q.shape[-1] in (64, 128) and falls back to pytorch_hstu_mha for other head_dims. If the CUTLASS forward kernel accepts an unsupported head_dim (e.g., 32) without raising — even if it silently misfires — the backward will compute gradients of pytorch_hstu_mha, not of the actual CUTLASS forward computation, violating the autograd.Function contract and producing wrong gradients. Adding the same guard in the forward ensures a clear ValueError rather than a cryptic kernel assertion and keeps forward/backward consistent.

Suggested change

	elif sm_major_version == 10:
	assert q.dtype in (
	torch.bfloat16,
	torch.float16,
	), f"Blackwell fwd expects bfloat16 or float16, got {q.dtype}"
	num_contexts = _blackwell_num_contexts_or_none(num_contexts)
	jagged_attn_output = hstu.hstu_attn_varlen_func(
	elif sm_major_version == 10:
	assert q.dtype in (
	torch.bfloat16,
	torch.float16,
	), f"Blackwell fwd expects bfloat16 or float16, got {q.dtype}"
	if q.shape[-1] not in (64, 128):
	raise ValueError(
	f"Blackwell fwd only supports head_dim in (64, 128), got {q.shape[-1]}"
	)
	num_contexts = _blackwell_num_contexts_or_none(num_contexts)
	jagged_attn_output = hstu.hstu_attn_varlen_func(

[JacoCheung]

pytorch-26.04 has triton bug. Need newer base image(26.05), which is not publicly available.

[JacoCheung]

Another constraint is blackwell does not support contextual features. Waiting for #395 .