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
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 newflashdreams.acceleratedPR (WIP)flashdreams.core.acceleration: CUDA graph, prewarm, and future flashdreams auto tune system. These feature are currently inflashdreams/infra/accelerationflashdreams.core.distributed: Distributed related features. Including context parallel (currently inflashdreams/core/attention/cp.pyandflashdreams/core/distributed), rank orchestration (currently inflashdreams/core/distributed)We also need to remove all the old
StreamInferencePipelinein favor of the composable inference pipeline once implemented. This should mean that we no longer needflashdreams/infraafter the refactor.Dependency
This issue is gated by #473 and #482