[lxd-cumt]

FlagOS Proposal: Plugin Architecture & Device-Agnostic Abstraction for TransformerEngine

Plugin System: Initialization-time Backend Loading

We propose a plugin architecture where TransformerEngine (TE) loads backend implementations at initialization time via an explicit plugin interface, while the actual multi-backend plugins reside in a separate repository (TransformerEngine-Plugin-FL).

TransformerEngine-Plugin-FL: https://github.com/lxd-cumt/TransformerEngine-Plugin-FL

Proposed Design

• TE defines a stable plugin API contract (operator signatures, quantization interfaces, communication primitives).
• At load_framework_extension() time, if a plugin is present, TE dispatches backend calls through the plugin registry; otherwise it falls back to the native te implementation.
• The TransformerEngine-Plugin-FL repository is independently installable and contains multiple backend implementations for diverse accelerators, and support more training scenarios.

[greptile-apps]

Confidence Score: 5/5

The changes since the previous review appear safe to merge, with only the existing non-blocking API-documentation concern remaining.

Summary

• Loads the module named by NVTE_PLUGIN during PyTorch extension initialization.
• Preserves the native extension under a _nv alias and rolls back module aliases when plugin initialization fails.
• Dispatches supported attention operations through plugin-provided implementations, with native fallbacks.
• The latest revision only adds the appropriate pylint suppression to the intentionally broad plugin-initialization exception handler.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Load native framework extension] --> B{Framework is torch and NVTE_PLUGIN is set?}
    B -->|No| C[Use native TransformerEngine extension]
    B -->|Yes| D[Register native extension under _nv alias]
    D --> E[Import configured plugin module]
    E --> F[Call load_plugins]
    F --> G{Plugin exposes attention overrides?}
    G -->|Yes| H[Install FlashAttention and backend-selection overrides]
    G -->|No| C
    E -->|Initialization fails| I[Remove alias and restore native module]
    F -->|Initialization fails| I
    I --> J[Emit RuntimeWarning]

Loading

[greptile-apps]

Document the plugin API contract

NVTE_PLUGIN introduces a user-facing configuration interface, but the required load_plugins(), flash_attention, and get_attention_backend contracts are not documented. Plugin implementers must infer their signatures and return values from source, leaving incompatible implementations to fail through import-time warnings or attention-time exceptions.

Knowledge Base Used:

• Common core (transformer_engine/common)
• PyTorch Attention Stack

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[ptrendx]

/te-ci pytorch

[ptrendx]

@lxd-cumt Please fix the lint issue (the one with too broad Exception catch).

[ptrendx]

/te-ci pytorch