Skip to content

[API] flashdreams.core Refactor #478

Description

@fangjunzhou-nv

Issue Description

Current flashdreams spreads the "core level" feature in multiple places under flashdreams/flashdreams. We need to refactor/reorganize the code layout to have a layered structure (flashdreams.core -> flashdreams.pipeline -> flashdreams.runtime)

Proposed Solution

The following features should go into the new flashdreams.core:

  • flashdreams.core.modules: accelerated flashdreams modules with their Triton kernel. Currently, they are in https://github.com/NVIDIA/flashdreams/tree/main/flashdreams/flashdreams/core/attention (special rope module & rope triton kernel) and my new flashdreams.accelerated PR (WIP)
  • flashdreams.core.acceleration: CUDA graph, prewarm, and future flashdreams auto tune system. These feature are currently in flashdreams/infra/acceleration
  • flashdreams.core.distributed: Distributed related features. Including context parallel (currently in flashdreams/core/attention/cp.py and flashdreams/core/distributed), rank orchestration (currently in flashdreams/core/distributed)

We also need to remove all the old StreamInferencePipeline in favor of the composable inference pipeline once implemented. This should mean that we no longer need flashdreams/infra after the refactor.

Dependency

This issue is gated by #473 and #482

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions