GLM5.3 Flash (KDA + mHC + KPool DSA) support + FP8

megatron/core/context_parallel_layout/__init__.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""Public context parallel sequence partition-mode APIs.
+
+The implementation is split by responsibility; internal conversion and route-building
+helpers remain in their respective submodules rather than being re-exported here.
+
+Ownership summary:
+
+- model builders choose the pipeline-stage input CP layout;
+- blocks convert rank-local sequence tensors between layer preferences;
+- model postprocess restores the public output boundary to the input layout;
+- MTP validates its inner-layer layout preference but does not own outer conversion.
+"""
+
+from megatron.core.context_parallel_layout.conversion import (
+    CpPartitionModeConverter,
+    convert_module_input_tensors_cp_partition_mode,
+)
+from megatron.core.context_parallel_layout.routes import prebuild_thd_cp_partition_routes
+from megatron.core.context_parallel_layout.types import CpPartitionMode, ThdCpRoute
+from megatron.core.context_parallel_layout.utils import finalize_packed_seq_params
+
+__all__ = [
+    "CpPartitionMode",
+    "CpPartitionModeConverter",
+    "ThdCpRoute",
+    "convert_module_input_tensors_cp_partition_mode",
+    "finalize_packed_seq_params",
+    "prebuild_thd_cp_partition_routes",
+]