Fix hstu inference cudagraph capture on sm100

[geoffreyQiu]

• Fix hstu inference cudagraph capture on sm100
• Make NV-Embedding-Cache as the repo submodule.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[geoffreyQiu]

/build devel=1

[JacoCheung]

✅ Pipeline #62096780 -- success

Job	Status	Log
pre_check	❌ failed	view
train_build_arm64	✅ success	view
prepare-jet-b200-smoke	✅ success	view
prepare-jet-b200-inference	✅ success	view
prepare-jet-cw-dfw-e2e-benchmark	✅ success	view
build_whl	✅ success	view
dynamicemb_test_fwd_bwd_8gpus	✅ success	view
dynamicemb_test_load_dump_8gpus	✅ success	view
unit_test_1gpu_a100	✅ success	view
unit_test_1gpu_h100	✅ success	view
unit_test_4gpu	✅ success	view
unit_test_tp_4gpu	❌ failed	view
L20_unit_test_1gpu	✅ success	view
inference_test_1gpu	✅ success	view
train_build_x86	✅ success	view
inference_unit_test_1gpu	✅ success	view

Result: 14/16 jobs passed

View full pipeline

[geoffreyQiu]

/build devel

[JacoCheung]

✅ Pipeline #62104884 -- success

Job	Status	Log
pre_check	❌ failed	view
train_build_x86	✅ success	view
train_build_arm64	✅ success	view
prepare-jet-b200-smoke	✅ success	view
prepare-jet-b200-inference	✅ success	view
prepare-jet-cw-dfw-e2e-benchmark	✅ success	view
build_whl	✅ success	view
dynamicemb_test_fwd_bwd_8gpus	✅ success	view
dynamicemb_test_load_dump_8gpus	✅ success	view
unit_test_1gpu_a100	✅ success	view
unit_test_4gpu	✅ success	view
unit_test_tp_4gpu	❌ failed	view
L20_unit_test_1gpu	✅ success	view
inference_unit_test_1gpu	✅ success	view
inference_test_1gpu	✅ success	view
unit_test_1gpu_h100	✅ success	view

Result: 14/16 jobs passed

View full pipeline

[JacoCheung]

/build devel

[JacoCheung]

/build devel

[JacoCheung]

/build devel

[JacoCheung]

✅ Pipeline #62272714 -- success

Job	Status	Log
build_devel_x86	✅ success	view
build_devel_arm64	✅ success	view
manifest_devel	✅ success	view
pre_check	❌ failed	view
train_build_x86	✅ success	view
train_build_arm64	✅ success	view
prepare-jet-b200-smoke	✅ success	view
prepare-jet-b200-inference	✅ success	view
prepare-jet-cw-dfw-e2e-benchmark	✅ success	view
build_whl	✅ success	view
dynamicemb_test_fwd_bwd_8gpus	✅ success	view
dynamicemb_test_load_dump_8gpus	✅ success	view
unit_test_1gpu_a100	✅ success	view
unit_test_1gpu_h100	✅ success	view
unit_test_4gpu	✅ success	view
unit_test_tp_4gpu	❌ failed	view
L20_unit_test_1gpu	✅ success	view
inference_unit_test_1gpu	✅ success	view
inference_test_1gpu	✅ success	view

Result: 17/19 jobs passed

View full pipeline

[greptile-apps]

Greptile Summary

The PR initializes valid synthetic paged-KV metadata before HSTU CUDA-graph capture on SM100 and vendors NV-Embedding-Cache as a pinned repository submodule.

• Adds capture-time KV page, offset, sequence-length, and token metadata initialization.
• Routes non-export paged attention through the FBGEMM HSTU function.
• Updates SM100 test cache sizing and Docker installation to use the pinned submodule.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
examples/hstu/modules/hstu_block_inference.py	Adds validated synthetic KV metadata initialization before CUDA-graph warmup and capture.
examples/hstu/modules/paged_hstu_infer_layer.py	Uses the FBGEMM paged-attention path during non-export inference.
examples/hstu/test/test_paged_hstu_layer.py	Adjusts SM100 page sizing and static metadata capacity for graph tests.
docker/Dockerfile	Builds NV-Embedding-Cache from the repository’s pinned submodule checkout.
third_party/nv-embedding-cache	Pins the NV-Embedding-Cache dependency as a Git submodule.

Sequence Diagram

sequenceDiagram
  participant Setup as CUDA graph setup
  participant Meta as Static KV metadata
  participant Layer as HSTU inference layer
  participant Attn as Paged HSTU attention
  Setup->>Meta: Initialize synthetic pages and offsets
  Setup->>Layer: Warm up and capture fixed shapes
  Layer->>Attn: Invoke paged attention
  Attn-->>Layer: Captured output
  Setup->>Meta: Copy request metadata
  Setup->>Layer: Replay matching graph

Loading

_{Reviews (2): Last reviewed commit: "Clean up code format" | Re-trigger Greptile}

[geoffreyQiu]

/build devel

[JacoCheung]

❌ Pipeline #62296888 -- failed

Job	Status	Log
build_devel_x86	❌ failed	view
build_devel_arm64	✅ success	view
manifest_devel	⏩ skipped	view
pre_check	❌ failed	view
train_build_x86	⏩ skipped	view
train_build_arm64	⏩ skipped	view
prepare-jet-b200-smoke	⏩ skipped	view
prepare-jet-b200-inference	⏩ skipped	view
prepare-jet-cw-dfw-e2e-benchmark	⏩ skipped	view
build_whl	⏩ skipped	view
dynamicemb_test_fwd_bwd_8gpus	⏩ skipped	view
dynamicemb_test_load_dump_8gpus	⏩ skipped	view
unit_test_1gpu_a100	⏩ skipped	view
unit_test_1gpu_h100	⏩ skipped	view
unit_test_4gpu	⏩ skipped	view
unit_test_tp_4gpu	⏩ skipped	view
L20_unit_test_1gpu	⏩ skipped	view
inference_unit_test_1gpu	⏩ skipped	view
inference_test_1gpu	⏩ skipped	view

Result: 1/3 jobs passed

View full pipeline