Fea transformer hstu inference 456

[Junyi-Zheng]

Description

Related to #456.

Add an opt-in Transformer backbone to the existing HSTU recommendation inference workflow, while keeping HSTU as the default.

• Implement Pre-LN softmax attention and a GELU FFN with causal history and independent candidate attention.
• Reuse recommendation preprocessing, positional embeddings, postprocessing, and paged KV-cache interfaces.
• Separate Transformer checkpoint loading from HSTU-specific weight transformations.
• Propagate backbone selection through ranking inference, Triton, and both AOTI export entry points.
• Add numerical, cache-consistency, checkpoint, and export regression tests, plus configuration and usage documentation.

Validation

• CPU regression tests: 132 passed, 3 skipped.
• Production-layer CPU AOTInductor compilation and replay passed.
• isort, autoflake, codespell, and black passed for changed files; the configured mypy check passed.
• git diff --check passed.

Limitations

CUDA-only tests were skipped because the local environment has no NVIDIA GPU. Native cache transfers, CUDA graph execution, complete GPU export, and Triton serving still require validation.

This is a reference padded-SDPA backend; GPU memory and performance have not been benchmarked. It requires its documented Transformer checkpoint layout.

Checklist

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

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; no concrete correctness, compatibility, security, or repository-rule issue remains.

Summary

This PR adds an opt-in Transformer backbone to the HSTU ranking-inference stack while preserving HSTU as the default.

• Implements Pre-LN multi-head attention, GELU feed-forward layers, recommendation-specific causal/candidate masking, and paged KV-cache integration.
• Propagates backbone selection through ranking inference, Triton, and both AOTI export workflows.
• Separates Transformer checkpoint loading from legacy HSTU weight conversion.
• Adds configuration validation, documentation, and numerical, cache, checkpoint, export, and CUDA integration tests.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Recommendation batch] --> B[Sparse embeddings]
  B --> C[Shared HSTU preprocessor]
  C --> D{NetworkArgs.backbone}
  D -->|hstu default| E[Paged HSTU layers]
  D -->|transformer opt-in| F[Pre-LN Transformer layers]
  G[Paged KV-cache metadata] --> E
  G --> F
  E --> H[Shared candidate postprocessor]
  F --> H
  H --> I[Shared prediction MLP]
  I --> J[Ranking outputs]
  D --> K[Python and Triton inference]
  D --> L[AOTI no-cache export]
  D --> M[AOTI KV-cache export]

Loading

_{Reviews (1) · Last reviewed commit: "Document Transformer interfaces and mark..."}