[feat] ULTRA-HSTU: FP8 attention via fbgemm_gpu_hstu wheel

[tiankongdeguiji]

Adds FP8 attention support to the ULTRA-HSTU stack. The fbgemm_gpu_hstu wheel already quantizes q/k/v internally (forward and backward) when hstu_attn_varlen_func is called with quant_mode>=0, so this PR is a thin pass-through: one int proto field on STU plumbed down to the CUTLASS attention call, plus an opaque export op so the FP8 path survives torch.export / AOTInductor.

Design

• Proto: optional int32 STU.fp8_quant_mode = 17 [default = -1] — -1 = off (bf16/fp16); 0..5 select an FP8 mode (per-tensor, two-direction, per-block, per-head, per-batch, global). The int mirrors the wheel's contract exactly and follows the existing scaling_seqlen=-1 / contextual_seq_len=-1 sentinel style. Lives on STU so it flows through config_to_kwargs(ch) → STULayer(**stu) with no HSTUTransducer.__init__ change, and each MoT channel can set its own.
• Threading: STULayer.__init__/forward → hstu_preprocess_and_attention → hstu_mha → cutlass_hstu_mha → hstu_attn_varlen_func(quant_mode=…). q/k/v stay bf16/fp16; the wheel does all FP8 quantization internally. No re-port of quantize helpers, no schema extension on an in-house op.
• Arch gate (fail-loud, no silent degradation): FP8 is accepted on SM90 (Hopper, all modes, fwd+bwd) or SM120 (Blackwell RTX, quant_mode=2 only, forward-only) — matches the wheel's per-arch dispatch. SM80/SM100 have no FP8 kernel and are rejected; SM120 quant_mode≠2 is rejected (wheel silently falls back to bf16 there). _assert_fp8_capable is is_fx_tracing-guarded so export on a non-capable box doesn't fault at trace time. hstu_mha rejects FP8 with any non-CUTLASS kernel; the fused-Triton preprocess path rejects it too. The cached/delta serving path is unchanged (bf16/fp16).
• Export / AOTInductor: the wheel quantizes q/k/v in Python (quantize_for_block_scale etc.) before the registered fbgemm::hstu_varlen_fwd_90 op. That quantizer iterates over per-sample lengths from cu_seqlens and emits descale tensors with a data-dependent block-count dim. torch.export traces straight through it (@torch.fx.wrap is symbolic-trace-only; dynamo/export ignore it), so modes 2/3/4/5 raise GuardOnDataDependentSymNode. Fix: register tzrec::cutlass_hstu_fp8_fwd as a torch.library.custom_op wrapping the entire FP8 forward (quantization + kernel) with a fake returning v.new_empty(v.shape). The exported graph then holds a single opaque op with backed shapes — no quantizer Python, no escaping descales. cutlass_hstu_mha routes to it only under torch.compiler.is_exporting(); eager train/eval keep calling hstu_attn_varlen_func directly (autograd intact), and predict runs the real quantizer inside the op via the AOTI proxy executor. Works with the stock wheel.
• Dep: bump fbgemm_gpu_hstu to 0.1.0+20260528.5f13f139.cu129. Version bump to 1.2.16.

Verification (H20, SM90)

• Unit (test_attn_fp8_cutlass, @mark_ci_scope("h20")): FP8 fwd + bwd vs bf16 PyTorch reference within atol/rtol≈0.1, sampling fp8_quant_mode ∈ {0, 2, 3, 4, 5} (mode 1 excluded — wheel-side vt TMA descriptor init fails for some shapes). SM90-gated, auto-skips on the local A10.
• Integration (test_rank_ultra_hstu_cutlass_train_eval_export): fp8_quant_mode=2 end-to-end — train + eval + AOTI export + predict all pass on H20. This config also has SLA (sla_k1/sla_k2) + attention truncation + MoT, so FP8 is exercised together with the NFUNC mask path through export and the compiled binary.
• Non-FP8 regression: fp8_quant_mode=-1 is bit-identical to omitting the kwarg, so the existing CUTLASS path is unchanged.

Known limitations

• FP8 quant_mode=1 is excluded from the unit parity test due to a wheel-side vt TMA-descriptor init failure on certain shapes; the tzrec code path supports it, but it isn't covered by a test.

User-facing docs land later in the consolidated ultra_hstu.md; this PR keeps docs/source/models/dlrm_hstu.md to the install-line update only, per the ULTRA-HSTU sub-PR convention.

🤖 Generated with Claude Code